diffusers 源码解析（三十七）
.\diffusers\pipelines\lumina\pipeline_lumina.py
# 版权所有 2024 Alpha-VLLM 和 HuggingFace 团队。保留所有权利。
#
# 根据 Apache 许可证，版本 2.0（“许可证”）进行许可；
# 您只能在遵守许可证的情况下使用此文件。
# 您可以在以下地址获取许可证的副本：
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律或书面协议另有约定，分发的软件均按“原样”提供，
# 不附带任何形式的明示或暗示的担保或条件。
# 有关许可证的具体权限和限制，请参见许可证。
import html # 导入处理 HTML 实体的模块
import inspect # 导入用于获取对象信息的模块
import math # 导入数学函数模块
import re # 导入正则表达式模块
import urllib.parse as ul # 导入用于 URL 解析的模块，并将其重命名为 ul
from typing import List, Optional, Tuple, Union # 导入类型注解模块中的相关类型
import torch # 导入 PyTorch 库
from transformers import AutoModel, AutoTokenizer # 从 transformers 导入自动模型和自动分词器
from ...image_processor import VaeImageProcessor # 从上级模块导入 VAE 图像处理器
from ...models import AutoencoderKL # 从上级模块导入自动编码器模型
from ...models.embeddings import get_2d_rotary_pos_embed_lumina # 从上级模块导入获取 2D 旋转位置嵌入的函数
from ...models.transformers.lumina_nextdit2d import LuminaNextDiT2DModel # 从上级模块导入 LuminaNextDiT2D 模型
from ...schedulers import FlowMatchEulerDiscreteScheduler # 从上级模块导入调度器
from ...utils import ( # 从上级模块导入多个实用工具
BACKENDS_MAPPING, # 后端映射
is_bs4_available, # 检查 bs4 库是否可用的函数
is_ftfy_available, # 检查 ftfy 库是否可用的函数
logging, # 日志记录模块
replace_example_docstring, # 替换示例文档字符串的函数
)
from ...utils.torch_utils import randn_tensor # 从上级模块导入生成随机张量的函数
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput # 从上级模块导入扩散管道和图像输出类
logger = logging.get_logger(__name__) # 获取当前模块的日志记录器
if is_bs4_available(): # 如果 bs4 库可用
from bs4 import BeautifulSoup # 导入 BeautifulSoup 库用于解析 HTML
if is_ftfy_available(): # 如果 ftfy 库可用
import ftfy # 导入 ftfy 库用于文本修复
EXAMPLE_DOC_STRING = """ # 示例文档字符串，用于展示用法
Examples:
```py
>>> import torch # 导入 PyTorch 库
>>> from diffusers import LuminaText2ImgPipeline # 从 diffusers 导入 LuminaText2ImgPipeline
>>> pipe = LuminaText2ImgPipeline.from_pretrained( # 从预训练模型加载管道
... "Alpha-VLLM/Lumina-Next-SFT-diffusers", torch_dtype=torch.bfloat16 # 指定模型路径和数据类型
... ).cuda() # 将管道移到 GPU
>>> # 启用内存优化。
>>> pipe.enable_model_cpu_offload() # 启用模型的 CPU 卸载以节省内存
>>> prompt = "Upper body of a young woman in a Victorian-era outfit with brass goggles and leather straps. Background shows an industrial revolution cityscape with smoky skies and tall, metal structures" # 定义生成图像的提示
>>> image = pipe(prompt).images[0] # 生成图像并提取第一张图像
```py
"""
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion 导入的函数
def retrieve_timesteps( # 定义函数以检索时间步
scheduler, # 调度器对象
num_inference_steps: Optional[int] = None, # 可选推理步骤数
device: Optional[Union[str, torch.device]] = None, # 可选设备类型
timesteps: Optional[List[int]] = None, # 可选时间步列表
sigmas: Optional[List[float]] = None, # 可选 sigma 值列表
**kwargs, # 接收其他关键字参数
):
"""
调用调度器的 `set_timesteps` 方法，并在调用后从调度器检索时间步。处理自定义时间步。
任何关键字参数将传递给 `scheduler.set_timesteps`。
# 参数说明
Args:
scheduler (`SchedulerMixin`): # 定义调度器类型，用于获取时间步长
The scheduler to get timesteps from. # 调度器的描述
num_inference_steps (`int`): # 定义推理步骤数量
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` # 生成样本时的扩散步骤数，如果使用，`timesteps` 必须为 `None`
must be `None`. # 说明条件
device (`str` or `torch.device`, *optional*): # 定义设备类型（字符串或torch.device），可选
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. # 指定将时间步移动到的设备，如果为 `None`，则不移动
timesteps (`List[int]`, *optional*): # 自定义时间步，列表类型，可选
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, # 自定义时间步，覆盖调度器的时间步策略，如果传递了 `timesteps`
`num_inference_steps` and `sigmas` must be `None`. # 则 `num_inference_steps` 和 `sigmas` 必须为 `None`
sigmas (`List[float]`, *optional*): # 自定义sigmas，列表类型，可选
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, # 自定义sigmas，覆盖调度器的时间步策略，如果传递了 `sigmas`
`num_inference_steps` and `timesteps` must be `None`. # 则 `num_inference_steps` 和 `timesteps` 必须为 `None`
Returns:
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the # 返回类型为元组，包含调度器的时间步计划和推理步骤数量
second element is the number of inference steps. # 返回第二个元素为推理步骤数量
"""
if timesteps is not None and sigmas is not None: # 检查是否同时传入了时间步和sigmas
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") # 抛出错误，要求只能选择一个
if timesteps is not None: # 如果传入了时间步
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) # 检查调度器是否接受时间步
if not accepts_timesteps: # 如果不接受
raise ValueError( # 抛出错误，说明当前调度器不支持自定义时间步
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" # 当前调度器不支持自定义时间步
f" timestep schedules. Please check whether you are using the correct scheduler." # 检查是否使用了正确的调度器
)
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) # 设置时间步，传入设备和其他参数
timesteps = scheduler.timesteps # 从调度器获取设置后的时间步
num_inference_steps = len(timesteps) # 计算推理步骤数量
elif sigmas is not None: # 如果传入了sigmas
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) # 检查调度器是否接受sigmas
if not accept_sigmas: # 如果不接受
raise ValueError( # 抛出错误，说明当前调度器不支持自定义sigmas
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" # 当前调度器不支持自定义sigmas
f" sigmas schedules. Please check whether you are using the correct scheduler." # 检查是否使用了正确的调度器
)
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) # 设置sigmas，传入设备和其他参数
timesteps = scheduler.timesteps # 从调度器获取设置后的时间步
num_inference_steps = len(timesteps) # 计算推理步骤数量
else: # 如果没有传入时间步和sigmas
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) # 直接使用推理步骤设置时间步，传入设备和其他参数
timesteps = scheduler.timesteps # 从调度器获取设置后的时间步
return timesteps, num_inference_steps # 返回时间步和推理步骤数量
# 定义 LuminaText2ImgPipeline 类，继承自 DiffusionPipeline
class LuminaText2ImgPipeline(DiffusionPipeline):
r"""
Lumina-T2I 的文本到图像生成管道。
此模型继承自 [`DiffusionPipeline`]。请查看父类文档以了解库为所有管道实现的通用方法（例如下载或保存、在特定设备上运行等）。
参数：
vae ([`AutoencoderKL`]):
用于编码和解码图像到潜在表示的变分自编码器（VAE）模型。
text_encoder ([`AutoModel`]):
冻结的文本编码器。Lumina-T2I 使用
[T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.AutoModel)，具体是
[t5-v1_1-xxl](https://huggingface.co/Alpha-VLLM/tree/main/t5-v1_1-xxl) 变体。
tokenizer (`AutoModel`):
[AutoModel](https://huggingface.co/docs/transformers/model_doc/t5#transformers.AutoModel) 类的分词器。
transformer ([`Transformer2DModel`]):
一种文本条件的 `Transformer2DModel`，用于去噪编码的图像潜在表示。
scheduler ([`SchedulerMixin`]):
用于与 `transformer` 结合使用的调度器，以去噪编码的图像潜在表示。
"""
# 编译一个正则表达式，用于匹配不良标点
bad_punct_regex = re.compile(
r"["
+ "#®•©™&@·º½¾¿¡§~"
+ r"\)"
+ r"\("
+ r"\]"
+ r"\["
+ r"\}"
+ r"\{"
+ r"\|"
+ "\\"
+ r"\/"
+ r"\*"
+ r"]{1,}"
) # noqa
# 定义可选组件的空列表
_optional_components = []
# 定义模型的 CPU 卸载顺序
model_cpu_offload_seq = "text_encoder->transformer->vae"
# 初始化方法，接受多个模型组件作为参数
def __init__(
self,
transformer: LuminaNextDiT2DModel,
scheduler: FlowMatchEulerDiscreteScheduler,
vae: AutoencoderKL,
text_encoder: AutoModel,
tokenizer: AutoTokenizer,
):
# 调用父类的初始化方法
super().__init__()
# 注册各个模块，便于后续使用
self.register_modules(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer,
scheduler=scheduler,
)
# 设置 VAE 的缩放因子
self.vae_scale_factor = 8
# 创建图像处理器，用于处理 VAE 输出
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
# 设置最大序列长度
self.max_sequence_length = 256
# 设置默认样本大小，根据 transformer 的配置或设置为 128
self.default_sample_size = (
self.transformer.config.sample_size
if hasattr(self, "transformer") and self.transformer is not None
else 128
)
# 计算默认图像大小
self.default_image_size = self.default_sample_size * self.vae_scale_factor
# 定义获取 gemma 提示嵌入的方法，接受多个参数
def _get_gemma_prompt_embeds(
self,
prompt: Union[str, List[str]],
num_images_per_prompt: int = 1,
device: Optional[torch.device] = None,
clean_caption: Optional[bool] = False,
max_length: Optional[int] = None,
# 处理输入的设备，使用给定设备或默认执行设备
):
device = device or self._execution_device
# 将字符串类型的提示转为列表形式
prompt = [prompt] if isinstance(prompt, str) else prompt
# 获取提示的批次大小
batch_size = len(prompt)
# 对提示文本进行预处理，选项可清理标题
prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
# 使用 tokenizer 处理提示文本，设置填充、最大长度等参数
text_inputs = self.tokenizer(
prompt,
pad_to_multiple_of=8,
max_length=self.max_sequence_length,
truncation=True,
padding=True,
return_tensors="pt",
)
# 将输入 ID 移动到指定设备
text_input_ids = text_inputs.input_ids.to(device)
# 获取未截断的输入 ID，使用最长填充
untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids.to(device)
# 检查是否存在截断，且文本输入 ID 与未截断 ID 不同
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
# 解码被截断的文本，并记录警告信息
removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.max_sequence_length - 1 : -1])
logger.warning(
"The following part of your input was truncated because Gemma can only handle sequences up to"
f" {self.max_sequence_length} tokens: {removed_text}"
)
# 将提示的注意力掩码移动到指定设备
prompt_attention_mask = text_inputs.attention_mask.to(device)
# 使用文本编码器获取提示的嵌入，输出隐藏状态
prompt_embeds = self.text_encoder(
text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True
)
# 获取倒数第二层的隐藏状态作为提示嵌入
prompt_embeds = prompt_embeds.hidden_states[-2]
# 确定数据类型，优先使用文本编码器的数据类型
if self.text_encoder is not None:
dtype = self.text_encoder.dtype
elif self.transformer is not None:
dtype = self.transformer.dtype
else:
dtype = None
# 将嵌入移动到指定的数据类型和设备
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
# 解包嵌入的形状以获取序列长度
_, seq_len, _ = prompt_embeds.shape
# 为每个提示生成多个图像，重复嵌入和注意力掩码
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
# 重新调整嵌入的形状
prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
# 重复注意力掩码以适应生成的图像数量
prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
# 重新调整注意力掩码的形状
prompt_attention_mask = prompt_attention_mask.view(batch_size * num_images_per_prompt, -1)
# 返回嵌入和注意力掩码
return prompt_embeds, prompt_attention_mask
# 从 diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt 适配
def encode_prompt(
# 定义编码提示的函数参数，包括提示和其他可选参数
prompt: Union[str, List[str]],
do_classifier_free_guidance: bool = True,
negative_prompt: Union[str, List[str]] = None,
num_images_per_prompt: int = 1,
device: Optional[torch.device] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
prompt_attention_mask: Optional[torch.Tensor] = None,
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
clean_caption: bool = False,
**kwargs,
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs 复制
# 定义一个方法，准备额外的参数以供调度器步骤使用
def prepare_extra_step_kwargs(self, generator, eta):
# 为调度器步骤准备额外的关键字参数，因为不同调度器的签名可能不同
# eta (η) 仅在 DDIMScheduler 中使用，其他调度器将忽略它
# eta 对应于 DDIM 论文中的 η: https://arxiv.org/abs/2010.02502
# eta 的值应在 [0, 1] 之间
# 检查调度器的 step 方法是否接受 eta 参数
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 初始化一个空的字典以存储额外的步骤参数
extra_step_kwargs = {}
# 如果调度器接受 eta 参数，将其添加到字典中
if accepts_eta:
extra_step_kwargs["eta"] = eta
# 检查调度器的 step 方法是否接受 generator 参数
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 如果调度器接受 generator 参数，将其添加到字典中
if accepts_generator:
extra_step_kwargs["generator"] = generator
# 返回包含额外参数的字典
return extra_step_kwargs
# 定义一个方法，用于检查输入的有效性
def check_inputs(
self,
prompt, # 输入的提示文本
height, # 图像的高度
width, # 图像的宽度
negative_prompt, # 输入的负面提示文本
prompt_embeds=None, # 可选的提示嵌入
negative_prompt_embeds=None, # 可选的负面提示嵌入
prompt_attention_mask=None, # 可选的提示注意力掩码
negative_prompt_attention_mask=None, # 可选的负面提示注意力掩码
):
# 检查高度和宽度是否是8的倍数，若不是，则抛出值错误
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
# 检查是否同时提供了 `prompt` 和 `prompt_embeds`，若是，则抛出值错误
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
# 检查是否都没有提供 `prompt` 和 `prompt_embeds`，若是，则抛出值错误
elif prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
# 检查 `prompt` 是否为字符串或列表，若不是，则抛出值错误
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
# 检查是否同时提供了 `prompt` 和 `negative_prompt_embeds`，若是，则抛出值错误
if prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
# 检查是否同时提供了 `negative_prompt` 和 `negative_prompt_embeds`，若是，则抛出值错误
if negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
# 检查如果提供了 `prompt_embeds`，则必须提供 `prompt_attention_mask`，若没有，则抛出值错误
if prompt_embeds is not None and prompt_attention_mask is None:
raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
# 检查如果提供了 `negative_prompt_embeds`，则必须提供 `negative_prompt_attention_mask`，若没有，则抛出值错误
if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
# 检查 `prompt_embeds` 和 `negative_prompt_embeds` 的形状是否一致，若不一致，则抛出值错误
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
# 检查 `prompt_attention_mask` 和 `negative_prompt_attention_mask` 的形状是否一致，若不一致，则抛出值错误
if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
raise ValueError(
"`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
f" {negative_prompt_attention_mask.shape}."
)
# 从 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing 复制的内容
# 定义文本预处理的私有方法，接受文本和一个可选的清理标题参数
def _text_preprocessing(self, text, clean_caption=False):
# 如果设置了清理标题且 bs4 库不可用，记录警告并将清理标题设为 False
if clean_caption and not is_bs4_available():
logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
logger.warning("Setting `clean_caption` to False...")
clean_caption = False
# 如果设置了清理标题且 ftfy 库不可用，记录警告并将清理标题设为 False
if clean_caption and not is_ftfy_available():
logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
logger.warning("Setting `clean_caption` to False...")
clean_caption = False
# 如果输入文本不是元组或列表，则将其转换为列表
if not isinstance(text, (tuple, list)):
text = [text]
# 定义内部处理文本的函数
def process(text: str):
# 如果设置了清理标题，则进行标题清理两次
if clean_caption:
text = self._clean_caption(text)
text = self._clean_caption(text)
else:
# 否则将文本转换为小写并去除空格
text = text.lower().strip()
return text
# 对文本列表中的每个元素进行处理，并返回处理后的结果列表
return [process(t) for t in text]
# 从 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption 复制而来
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
# 定义潜在变量的形状，根据输入参数计算
shape = (
batch_size,
num_channels_latents,
int(height) // self.vae_scale_factor,
int(width) // self.vae_scale_factor,
)
# 如果生成器是列表且长度与批量大小不匹配，则引发错误
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
# 如果潜在变量为空，则生成随机张量作为潜在变量
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 如果潜在变量已给出，则将其转移到指定设备
latents = latents.to(device)
# 返回生成的或已转换的潜在变量
return latents
# 返回指导比例的属性值
@property
def guidance_scale(self):
return self._guidance_scale
# 此处的指导比例与 Imagen 论文中方程 (2) 的指导权重 `w` 相似： https://arxiv.org/pdf/2205.11487.pdf
# 指导比例 = 1 表示没有进行无分类器的指导
@property
def do_classifier_free_guidance(self):
return self._guidance_scale > 1
# 返回时间步数的属性值
@property
def num_timesteps(self):
return self._num_timesteps
# 关闭梯度计算以节省内存
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义可调用对象的方法
def __call__(
# 输入的提示，可以是字符串或字符串列表
self,
prompt: Union[str, List[str]] = None,
# 输出图像的宽度
width: Optional[int] = None,
# 输出图像的高度
height: Optional[int] = None,
# 推理步骤的数量，默认为30
num_inference_steps: int = 30,
# 定义时间步的列表，默认为None
timesteps: List[int] = None,
# 引导尺度，默认为4.0
guidance_scale: float = 4.0,
# 负提示，可以是字符串或字符串列表
negative_prompt: Union[str, List[str]] = None,
# sigma值的列表，默认为None
sigmas: List[float] = None,
# 每个提示生成的图像数量，默认为1
num_images_per_prompt: Optional[int] = 1,
# 随机数生成器，可以是单个或多个torch生成器
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 潜在变量，默认为None
latents: Optional[torch.Tensor] = None,
# 提示的嵌入向量，默认为None
prompt_embeds: Optional[torch.Tensor] = None,
# 负提示的嵌入向量，默认为None
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 提示的注意力掩码，默认为None
prompt_attention_mask: Optional[torch.Tensor] = None,
# 负提示的注意力掩码，默认为None
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
# 输出类型，默认为"pil"
output_type: Optional[str] = "pil",
# 是否返回字典格式，默认为True
return_dict: bool = True,
# 是否清理提示文本，默认为True
clean_caption: bool = True,
# 最大序列长度，默认为256
max_sequence_length: int = 256,
# 缩放阈值，默认为1.0
scaling_watershed: Optional[float] = 1.0,
# 是否使用比例注意力，默认为True
proportional_attn: Optional[bool] = True,
.\diffusers\pipelines\lumina\__init__.py
# 导入类型检查模块
from typing import TYPE_CHECKING
# 从相对路径的 utils 模块中导入多个工具和常量
from ...utils import (
DIFFUSERS_SLOW_IMPORT, # 表示是否慢加载
OptionalDependencyNotAvailable, # 可选依赖未找到异常
_LazyModule, # 延迟加载模块的工具
get_objects_from_module, # 从模块中获取对象的工具
is_torch_available, # 检查 PyTorch 是否可用的工具
is_transformers_available, # 检查 Transformers 是否可用的工具
)
# 存储虚拟对象的字典
_dummy_objects = {}
# 存储模块导入结构的字典
_import_structure = {}
# 尝试执行以下代码块
try:
# 检查 Transformers 和 PyTorch 是否可用，若都不可用则抛出异常
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable()
# 捕获可选依赖未找到异常
except OptionalDependencyNotAvailable:
# 从 utils 模块中导入虚拟对象（占位符）
from ...utils import dummy_torch_and_transformers_objects # noqa F403
# 更新 _dummy_objects 字典，加入虚拟对象
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
# 如果没有异常，执行以下代码
else:
# 在导入结构字典中添加键值对，表示可用的管道
_import_structure["pipeline_lumina"] = ["LuminaText2ImgPipeline"]
# 如果是类型检查或者慢加载标志为真，执行以下代码块
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
try:
# 再次检查 Transformers 和 PyTorch 是否可用，若都不可用则抛出异常
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable()
# 捕获可选依赖未找到异常
except OptionalDependencyNotAvailable:
# 从 utils 中导入虚拟对象（占位符）
from ...utils.dummy_torch_and_transformers_objects import *
else:
# 从 pipeline_lumina 模块导入 LuminaText2ImgPipeline
from .pipeline_lumina import LuminaText2ImgPipeline
# 否则，执行以下代码块
else:
# 导入 sys 模块
import sys
# 将当前模块替换为延迟加载模块
sys.modules[__name__] = _LazyModule(
__name__, # 模块名称
globals()["__file__"], # 当前文件路径
_import_structure, # 导入结构
module_spec=__spec__, # 模块规格
)
# 遍历 _dummy_objects 字典，将每个虚拟对象设置到当前模块
for name, value in _dummy_objects.items():
setattr(sys.modules[__name__], name, value)
.\diffusers\pipelines\marigold\marigold_image_processing.py
# 从类型提示模块导入所需类型
from typing import List, Optional, Tuple, Union
# 导入 NumPy 库作为 np
import numpy as np
# 导入 PIL 库
import PIL
# 导入 PyTorch 库
import torch
# 导入 PyTorch 的功能模块
import torch.nn.functional as F
# 从 PIL 导入 Image 类
from PIL import Image
# 从上级模块导入 ConfigMixin 类
from ... import ConfigMixin
# 从配置工具模块导入注册配置的装饰器
from ...configuration_utils import register_to_config
# 从图像处理模块导入管道图像输入类
from ...image_processor import PipelineImageInput
# 从工具模块导入配置名称和日志工具
from ...utils import CONFIG_NAME, logging
# 从导入工具模块导入判断是否可用的 Matplotlib
from ...utils.import_utils import is_matplotlib_available
# 创建一个日志记录器，使用当前模块的名称
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 定义 MarigoldImageProcessor 类，继承自 ConfigMixin
class MarigoldImageProcessor(ConfigMixin):
# 设置配置名称为常量 CONFIG_NAME
config_name = CONFIG_NAME
# 注册初始化方法到配置中
@register_to_config
def __init__(
self,
vae_scale_factor: int = 8, # 变分自编码器的缩放因子，默认为 8
do_normalize: bool = True, # 是否进行归一化，默认为 True
do_range_check: bool = True, # 是否进行范围检查，默认为 True
):
super().__init__() # 调用父类的初始化方法
# 定义静态方法，扩展张量或数组
@staticmethod
def expand_tensor_or_array(images: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]:
"""
扩展张量或数组到指定数量的图像。
"""
if isinstance(images, np.ndarray): # 检查是否为 NumPy 数组
if images.ndim == 2: # 如果是二维数组 [H,W] -> [1,H,W,1]
images = images[None, ..., None]
if images.ndim == 3: # 如果是三维数组 [H,W,C] -> [1,H,W,C]
images = images[None]
elif isinstance(images, torch.Tensor): # 检查是否为 PyTorch 张量
if images.ndim == 2: # 如果是二维张量 [H,W] -> [1,1,H,W]
images = images[None, None]
elif images.ndim == 3: # 如果是三维张量 [1,H,W] -> [1,1,H,W]
images = images[None]
else: # 如果类型不匹配，抛出错误
raise ValueError(f"Unexpected input type: {type(images)}")
return images # 返回处理后的图像
# 定义静态方法，将 PyTorch 张量转换为 NumPy 图像
@staticmethod
def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
"""
将 PyTorch 张量转换为 NumPy 图像。
"""
images = images.cpu().permute(0, 2, 3, 1).float().numpy() # 转移到 CPU，调整维度并转换为 NumPy 数组
return images # 返回转换后的图像
# 定义静态方法，将 NumPy 图像转换为 PyTorch 张量
@staticmethod
def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
"""
将 NumPy 图像转换为 PyTorch 张量。
"""
if np.issubdtype(images.dtype, np.integer) and not np.issubdtype(images.dtype, np.unsignedinteger):
# 如果图像数据类型是有符号整数，抛出错误
raise ValueError(f"Input image dtype={images.dtype} cannot be a signed integer.")
if np.issubdtype(images.dtype, np.complexfloating):
# 如果图像数据类型是复数，抛出错误
raise ValueError(f"Input image dtype={images.dtype} cannot be complex.")
if np.issubdtype(images.dtype, bool):
# 如果图像数据类型是布尔型，抛出错误
raise ValueError(f"Input image dtype={images.dtype} cannot be boolean.")
images = torch.from_numpy(images.transpose(0, 3, 1, 2)) # 将 NumPy 数组转换为 PyTorch 张量，并调整维度
return images # 返回转换后的张量
# 定义静态方法，用于带抗锯齿的图像调整大小
@staticmethod
def resize_antialias(
image: torch.Tensor, size: Tuple[int, int], mode: str, is_aa: Optional[bool] = None # 接受图像、目标大小、模式和抗锯齿参数
) -> torch.Tensor: # 定义一个返回类型为 torch.Tensor 的函数
# 检查输入是否为 tensor，如果不是，抛出 ValueError 异常
if not torch.is_tensor(image):
raise ValueError(f"Invalid input type={type(image)}.")
# 检查输入的 dtype 是否为浮点型，如果不是，抛出 ValueError 异常
if not torch.is_floating_point(image):
raise ValueError(f"Invalid input dtype={image.dtype}.")
# 检查输入的维度是否为 4，如果不是，抛出 ValueError 异常
if image.dim() != 4:
raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
# 判断是否需要抗锯齿处理，并检查模式是否为双线性或双三次插值
antialias = is_aa and mode in ("bilinear", "bicubic")
# 对输入图像进行插值处理，调整到指定大小
image = F.interpolate(image, size, mode=mode, antialias=antialias)
# 返回调整大小后的图像
return image
@staticmethod # 表示这是一个静态方法
def resize_to_max_edge(image: torch.Tensor, max_edge_sz: int, mode: str) -> torch.Tensor: # 定义一个返回类型为 torch.Tensor 的函数
# 检查输入是否为 tensor，如果不是，抛出 ValueError 异常
if not torch.is_tensor(image):
raise ValueError(f"Invalid input type={type(image)}.")
# 检查输入的 dtype 是否为浮点型，如果不是，抛出 ValueError 异常
if not torch.is_floating_point(image):
raise ValueError(f"Invalid input dtype={image.dtype}.")
# 检查输入的维度是否为 4，如果不是，抛出 ValueError 异常
if image.dim() != 4:
raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
# 获取图像的高度和宽度
h, w = image.shape[-2:]
# 计算原始图像的最大边长度
max_orig = max(h, w)
# 计算新的高度和宽度，保持最大边不超过 max_edge_sz
new_h = h * max_edge_sz // max_orig
new_w = w * max_edge_sz // max_orig
# 检查新的高度和宽度是否为 0，如果是，抛出 ValueError 异常
if new_h == 0 or new_w == 0:
raise ValueError(f"Extreme aspect ratio of the input image: [{w} x {h}]")
# 调用静态方法进行抗锯齿处理并调整图像大小
image = MarigoldImageProcessor.resize_antialias(image, (new_h, new_w), mode, is_aa=True)
# 返回调整后的图像
return image
@staticmethod # 表示这是一个静态方法
def pad_image(image: torch.Tensor, align: int) -> Tuple[torch.Tensor, Tuple[int, int]]: # 定义一个返回类型为元组的函数
# 检查输入是否为 tensor，如果不是，抛出 ValueError 异常
if not torch.is_tensor(image):
raise ValueError(f"Invalid input type={type(image)}.")
# 检查输入的 dtype 是否为浮点型，如果不是，抛出 ValueError 异常
if not torch.is_floating_point(image):
raise ValueError(f"Invalid input dtype={image.dtype}.")
# 检查输入的维度是否为 4，如果不是，抛出 ValueError 异常
if image.dim() != 4:
raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
# 获取图像的高度和宽度
h, w = image.shape[-2:]
# 计算需要的填充量以满足对齐要求
ph, pw = -h % align, -w % align
# 使用重复模式对图像进行填充
image = F.pad(image, (0, pw, 0, ph), mode="replicate")
# 返回填充后的图像和填充量
return image, (ph, pw)
@staticmethod # 表示这是一个静态方法
def unpad_image(image: torch.Tensor, padding: Tuple[int, int]) -> torch.Tensor: # 定义一个返回类型为 torch.Tensor 的函数
# 检查输入是否为 tensor，如果不是，抛出 ValueError 异常
if not torch.is_tensor(image):
raise ValueError(f"Invalid input type={type(image)}.")
# 检查输入的 dtype 是否为浮点型，如果不是，抛出 ValueError 异常
if not torch.is_floating_point(image):
raise ValueError(f"Invalid input dtype={image.dtype}.")
# 检查输入的维度是否为 4，如果不是，抛出 ValueError 异常
if image.dim() != 4:
raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
# 从填充元组中提取高度和宽度的填充量
ph, pw = padding
# 如果填充量为 0，设置为 None，否则取负填充量
uh = None if ph == 0 else -ph
uw = None if pw == 0 else -pw
# 根据计算的新的高度和宽度裁剪图像
image = image[:, :, :uh, :uw]
# 返回裁剪后的图像
return image
@staticmethod # 表示这是一个静态方法
def load_image_canonical( # 定义一个接受多种类型的输入图像的函数
image: Union[torch.Tensor, np.ndarray, Image.Image], # 接受 tensor、numpy 数组或 PIL 图像
device: torch.device = torch.device("cpu"), # 设置默认设备为 CPU
dtype: torch.dtype = torch.float32, # 设置默认数据类型为 float32
# 返回类型为元组，包括一个张量和一个整数
) -> Tuple[torch.Tensor, int]:
# 检查输入是否为 PIL 图像类型
if isinstance(image, Image.Image):
# 将 PIL 图像转换为 NumPy 数组
image = np.array(image)
# 初始化图像数据类型的最大值
image_dtype_max = None
# 检查输入是否为 NumPy 数组或 PyTorch 张量
if isinstance(image, (np.ndarray, torch.Tensor)):
# 扩展张量或数组的维度
image = MarigoldImageProcessor.expand_tensor_or_array(image)
# 确保图像的维度是 2、3 或 4
if image.ndim != 4:
raise ValueError("Input image is not 2-, 3-, or 4-dimensional.")
# 检查输入是否为 NumPy 数组
if isinstance(image, np.ndarray):
# 检查图像数据类型是否为有符号整数
if np.issubdtype(image.dtype, np.integer) and not np.issubdtype(image.dtype, np.unsignedinteger):
raise ValueError(f"Input image dtype={image.dtype} cannot be a signed integer.")
# 检查图像数据类型是否为复数
if np.issubdtype(image.dtype, np.complexfloating):
raise ValueError(f"Input image dtype={image.dtype} cannot be complex.")
# 检查图像数据类型是否为布尔值
if np.issubdtype(image.dtype, bool):
raise ValueError(f"Input image dtype={image.dtype} cannot be boolean.")
# 检查图像数据类型是否为无符号整数
if np.issubdtype(image.dtype, np.unsignedinteger):
# 获取无符号整数的最大值
image_dtype_max = np.iinfo(image.dtype).max
# 转换数据类型为浮点数
image = image.astype(np.float32) # 因为 torch 不支持无符号数据类型超过 torch.uint8
# 将 NumPy 数组转换为 PyTorch 张量
image = MarigoldImageProcessor.numpy_to_pt(image)
# 检查是否为张量并且不是浮点类型且没有最大值
if torch.is_tensor(image) and not torch.is_floating_point(image) and image_dtype_max is None:
# 检查图像数据类型是否为 uint8
if image.dtype != torch.uint8:
raise ValueError(f"Image dtype={image.dtype} is not supported.")
# 设置最大值为 255
image_dtype_max = 255
# 确保输入是张量
if not torch.is_tensor(image):
raise ValueError(f"Input type unsupported: {type(image)}.")
# 如果图像的通道数为 1，则重复通道以形成 RGB 图像
if image.shape[1] == 1:
image = image.repeat(1, 3, 1, 1) # [N,1,H,W] -> [N,3,H,W]
# 确保图像是 1 通道或 3 通道
if image.shape[1] != 3:
raise ValueError(f"Input image is not 1- or 3-channel: {image.shape}.")
# 将图像移动到指定设备，并转换为指定数据类型
image = image.to(device=device, dtype=dtype)
# 如果存在数据类型最大值，则将图像数据归一化
if image_dtype_max is not None:
image = image / image_dtype_max
# 返回处理后的图像
return image
# 静态方法，检查图像的值范围
@staticmethod
def check_image_values_range(image: torch.Tensor) -> None:
# 确保输入是张量
if not torch.is_tensor(image):
raise ValueError(f"Invalid input type={type(image)}.")
# 确保输入是浮点类型
if not torch.is_floating_point(image):
raise ValueError(f"Invalid input dtype={image.dtype}.")
# 检查图像数据是否在 [0,1] 范围内
if image.min().item() < 0.0 or image.max().item() > 1.0:
raise ValueError("Input image data is partially outside of the [0,1] range.")
# 预处理方法
def preprocess(
self,
# 输入图像，可以是多种格式
image: PipelineImageInput,
# 可选的处理分辨率
processing_resolution: Optional[int] = None,
# 输入的重采样方法
resample_method_input: str = "bilinear",
# 指定设备（CPU 或 GPU）
device: torch.device = torch.device("cpu"),
# 指定数据类型，默认为浮点数
dtype: torch.dtype = torch.float32,
):
# 检查输入的图像是否为列表类型
if isinstance(image, list):
# 初始化图像变量
images = None
# 遍历图像列表，获取每个图像的索引和内容
for i, img in enumerate(image):
# 加载图像并标准化为指定的格式，返回形状为[N,3,H,W]
img = self.load_image_canonical(img, device, dtype) # [N,3,H,W]
# 如果还没有图像，直接赋值
if images is None:
images = img
else:
# 检查当前图像的维度是否与已有图像兼容
if images.shape[2:] != img.shape[2:]:
# 如果不兼容，抛出错误并给出详细信息
raise ValueError(
f"Input image[{i}] has incompatible dimensions {img.shape[2:]} with the previous images "
f"{images.shape[2:]}"
)
# 将当前图像与已有图像在第一维拼接
images = torch.cat((images, img), dim=0)
# 将最终图像集赋值回原变量
image = images
# 删除临时图像变量以释放内存
del images
else:
# 加载单个图像并标准化为指定的格式，返回形状为[N,3,H,W]
image = self.load_image_canonical(image, device, dtype) # [N,3,H,W]
# 获取图像的原始分辨率
original_resolution = image.shape[2:]
# 如果配置要求进行值范围检查，则执行检查
if self.config.do_range_check:
self.check_image_values_range(image)
# 如果配置要求进行归一化处理，则进行操作
if self.config.do_normalize:
image = image * 2.0 - 1.0
# 如果处理分辨率被指定且大于0，则调整图像大小
if processing_resolution is not None and processing_resolution > 0:
# 调整图像到最大边长，返回形状为[N,3,PH,PW]
image = self.resize_to_max_edge(image, processing_resolution, resample_method_input) # [N,3,PH,PW]
# 对图像进行填充，返回填充后的图像和填充信息，形状为[N,3,PPH,PPW]
image, padding = self.pad_image(image, self.config.vae_scale_factor) # [N,3,PPH,PPW]
# 返回处理后的图像、填充信息和原始分辨率
return image, padding, original_resolution
# 定义静态方法 colormap，用于图像上色
@staticmethod
def colormap(
# 输入图像，支持多种类型
image: Union[np.ndarray, torch.Tensor],
# 颜色映射名称，默认为 "Spectral"
cmap: str = "Spectral",
# 是否返回字节类型的图像
bytes: bool = False,
# 强制使用的特定方法，默认为 None
_force_method: Optional[str] = None,
# 定义静态方法 visualize_depth，用于可视化深度信息
@staticmethod
def visualize_depth(
# 输入深度图像，支持多种类型
depth: Union[
PIL.Image.Image,
np.ndarray,
torch.Tensor,
List[PIL.Image.Image],
List[np.ndarray],
List[torch.Tensor],
],
# 深度值的最小阈值，默认为0.0
val_min: float = 0.0,
# 深度值的最大阈值，默认为1.0
val_max: float = 1.0,
# 颜色映射名称，默认为 "Spectral"
color_map: str = "Spectral",
# 返回深度图像的可视化结果，可以是单个图像或图像列表
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
"""
可视化深度图，例如 `MarigoldDepthPipeline` 的预测结果。
参数:
depth (`Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image], List[np.ndarray],
List[torch.Tensor]]`): 深度图。
val_min (`float`, *可选*, 默认值为 `0.0`): 可视化深度范围的最小值。
val_max (`float`, *可选*, 默认值为 `1.0`): 可视化深度范围的最大值。
color_map (`str`, *可选*, 默认值为 `"Spectral"`): 用于将单通道深度预测转换为彩色表示的颜色映射。
返回: `PIL.Image.Image` 或 `List[PIL.Image.Image]`，包含深度图可视化结果。
"""
# 检查最大值是否小于等于最小值，若是则抛出错误
if val_max <= val_min:
raise ValueError(f"Invalid values range: [{val_min}, {val_max}].")
# 定义用于可视化单个深度图的函数
def visualize_depth_one(img, idx=None):
# 为图像前缀生成字符串，包含索引（如果存在）
prefix = "Depth" + (f"[{idx}]" if idx else "")
# 检查输入图像是否为 PIL 图像
if isinstance(img, PIL.Image.Image):
# 验证图像模式是否为 "I;16"
if img.mode != "I;16":
raise ValueError(f"{prefix}: invalid PIL mode={img.mode}.")
# 将 PIL 图像转换为 numpy 数组并归一化
img = np.array(img).astype(np.float32) / (2**16 - 1)
# 检查输入图像是否为 numpy 数组或 PyTorch 张量
if isinstance(img, np.ndarray) or torch.is_tensor(img):
# 确保输入图像是二维的
if img.ndim != 2:
raise ValueError(f"{prefix}: unexpected shape={img.shape}.")
# 若为 numpy 数组，则转换为 PyTorch 张量
if isinstance(img, np.ndarray):
img = torch.from_numpy(img)
# 确保图像是浮点类型
if not torch.is_floating_point(img):
raise ValueError(f"{prefix}: unexected dtype={img.dtype}.")
else:
# 如果输入类型不匹配，则抛出错误
raise ValueError(f"{prefix}: unexpected type={type(img)}.")
# 如果最小值或最大值不为默认值，则进行归一化处理
if val_min != 0.0 or val_max != 1.0:
img = (img - val_min) / (val_max - val_min)
# 使用颜色映射处理深度图像并转换为 RGB 格式
img = MarigoldImageProcessor.colormap(img, cmap=color_map, bytes=True) # [H,W,3]
# 将数组转换回 PIL 图像
img = PIL.Image.fromarray(img.cpu().numpy())
return img
# 检查输入深度是否为 None 或列表中的元素为 None
if depth is None or isinstance(depth, list) and any(o is None for o in depth):
raise ValueError("Input depth is `None`")
# 如果输入深度为 numpy 数组或 PyTorch 张量
if isinstance(depth, (np.ndarray, torch.Tensor)):
# 扩展张量或数组以匹配预期形状
depth = MarigoldImageProcessor.expand_tensor_or_array(depth)
# 若为 numpy 数组，则转换为 PyTorch 张量，形状调整为 [N,1,H,W]
if isinstance(depth, np.ndarray):
depth = MarigoldImageProcessor.numpy_to_pt(depth) # [N,H,W,1] -> [N,1,H,W]
# 验证深度图形状是否符合预期
if not (depth.ndim == 4 and depth.shape[1] == 1): # [N,1,H,W]
raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].")
# 返回每个图像的可视化结果列表
return [visualize_depth_one(img[0], idx) for idx, img in enumerate(depth)]
# 如果输入深度为列表，则对每个图像进行可视化
elif isinstance(depth, list):
return [visualize_depth_one(img, idx) for idx, img in enumerate(depth)]
else:
# 如果输入类型不匹配，则抛出错误
raise ValueError(f"Unexpected input type: {type(depth)}")
# 定义静态方法标识
@staticmethod
# 导出深度图为16位PNG格式
def export_depth_to_16bit_png(
# 深度数据，支持多种输入格式
depth: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]],
# 深度值的最小范围
val_min: float = 0.0,
# 深度值的最大范围
val_max: float = 1.0,
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
# 导出单张深度图为16位PNG格式的内部函数
def export_depth_to_16bit_png_one(img, idx=None):
# 生成深度图的前缀，用于错误信息
prefix = "Depth" + (f"[{idx}]" if idx else "")
# 检查输入是否为有效类型
if not isinstance(img, np.ndarray) and not torch.is_tensor(img):
raise ValueError(f"{prefix}: unexpected type={type(img)}.")
# 检查输入的维度是否为2D
if img.ndim != 2:
raise ValueError(f"{prefix}: unexpected shape={img.shape}.")
# 将PyTorch张量转换为NumPy数组
if torch.is_tensor(img):
img = img.cpu().numpy()
# 检查数据类型是否为浮点数
if not np.issubdtype(img.dtype, np.floating):
raise ValueError(f"{prefix}: unexected dtype={img.dtype}.")
# 根据给定范围标准化深度图
if val_min != 0.0 or val_max != 1.0:
img = (img - val_min) / (val_max - val_min)
# 将深度图值转换为16位整数
img = (img * (2**16 - 1)).astype(np.uint16)
# 将数组转换为16位PNG格式的图像
img = PIL.Image.fromarray(img, mode="I;16")
# 返回生成的图像
return img
# 检查输入深度数据是否为None或包含None
if depth is None or isinstance(depth, list) and any(o is None for o in depth):
raise ValueError("Input depth is `None`")
# 如果输入为NumPy数组或PyTorch张量
if isinstance(depth, (np.ndarray, torch.Tensor)):
# 扩展张量或数组的维度
depth = MarigoldImageProcessor.expand_tensor_or_array(depth)
# 如果输入是NumPy数组，转换为PyTorch张量
if isinstance(depth, np.ndarray):
depth = MarigoldImageProcessor.numpy_to_pt(depth) # [N,H,W,1] -> [N,1,H,W]
# 检查扩展后的深度图形状
if not (depth.ndim == 4 and depth.shape[1] == 1):
raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].")
# 返回每张深度图的16位PNG图像
return [export_depth_to_16bit_png_one(img[0], idx) for idx, img in enumerate(depth)]
# 如果输入是列表
elif isinstance(depth, list):
# 返回每张深度图的16位PNG图像
return [export_depth_to_16bit_png_one(img, idx) for idx, img in enumerate(depth)]
else:
# 抛出不支持的输入类型错误
raise ValueError(f"Unexpected input type: {type(depth)}")
# 可视化法线的静态方法
@staticmethod
def visualize_normals(
# 法线数据，支持多种输入格式
normals: Union[
np.ndarray,
torch.Tensor,
List[np.ndarray],
List[torch.Tensor],
],
# 是否沿X轴翻转
flip_x: bool = False,
# 是否沿Y轴翻转
flip_y: bool = False,
# 是否沿Z轴翻转
flip_z: bool = False,
# 返回类型为 PIL.Image.Image 或 List[PIL.Image.Image]，用于可视化表面法线
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
"""
可视化表面法线，例如 `MarigoldNormalsPipeline` 的预测结果。
参数：
normals (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
表面法线数据。
flip_x (`bool`, *可选*, 默认值为 `False`): 翻转法线参考系的 X 轴。
默认方向为右。
flip_y (`bool`, *可选*, 默认值为 `False`): 翻转法线参考系的 Y 轴。
默认方向为上。
flip_z (`bool`, *可选*, 默认值为 `False`): 翻转法线参考系的 Z 轴。
默认方向为面向观察者。
返回值: `PIL.Image.Image` 或 `List[PIL.Image.Image]`，包含表面法线的可视化图像。
"""
# 初始化翻转向量为 None
flip_vec = None
# 如果任一翻转标志为真，则创建翻转向量
if any((flip_x, flip_y, flip_z)):
flip_vec = torch.tensor(
[
(-1) ** flip_x, # 根据 flip_x 计算 X 轴的翻转因子
(-1) ** flip_y, # 根据 flip_y 计算 Y 轴的翻转因子
(-1) ** flip_z, # 根据 flip_z 计算 Z 轴的翻转因子
],
dtype=torch.float32, # 数据类型为浮点数
)
# 定义一个用于可视化单个法线图像的函数
def visualize_normals_one(img, idx=None):
img = img.permute(1, 2, 0) # 改变图像维度顺序为 (H, W, C)
if flip_vec is not None:
img *= flip_vec.to(img.device) # 应用翻转向量
img = (img + 1.0) * 0.5 # 将图像数据归一化到 [0, 1]
img = (img * 255).to(dtype=torch.uint8, device="cpu").numpy() # 转换为 uint8 格式并转为 numpy 数组
img = PIL.Image.fromarray(img) # 将 numpy 数组转换为 PIL 图像
return img # 返回处理后的图像
# 检查输入法线是否为 None 或含有 None 的列表
if normals is None or isinstance(normals, list) and any(o is None for o in normals):
raise ValueError("Input normals is `None`") # 抛出异常
# 如果法线数据为 numpy 数组或 torch 张量
if isinstance(normals, (np.ndarray, torch.Tensor)):
normals = MarigoldImageProcessor.expand_tensor_or_array(normals) # 扩展法线数据
if isinstance(normals, np.ndarray):
normals = MarigoldImageProcessor.numpy_to_pt(normals) # 转换 numpy 数组为 PyTorch 张量，形状为 [N,3,H,W]
# 检查法线数据的维度和形状
if not (normals.ndim == 4 and normals.shape[1] == 3):
raise ValueError(f"Unexpected input shape={normals.shape}, expecting [N,3,H,W].") # 抛出异常
# 可视化每个法线图像并返回图像列表
return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)]
# 如果法线数据为列表
elif isinstance(normals, list):
# 可视化每个法线图像并返回图像列表
return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)]
else:
raise ValueError(f"Unexpected input type: {type(normals)}") # 抛出异常，处理未知类型
# 定义静态方法可视化不确定性
@staticmethod
def visualize_uncertainty(
uncertainty: Union[
np.ndarray,
torch.Tensor,
List[np.ndarray],
List[torch.Tensor],
],
saturation_percentile=95, # 定义饱和度百分位参数，默认为95%
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
# 指定函数返回类型为单个 PIL.Image.Image 或者 PIL.Image.Image 的列表
"""
# 文档字符串，说明函数的功能，参数及返回值
Visualizes dense uncertainties, such as produced by `MarigoldDepthPipeline` or `MarigoldNormalsPipeline`.
Args:
# 参数说明，uncertainty 可以是不同类型的数组
uncertainty (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
Uncertainty maps.
# 参数说明，饱和度百分位数，默认为 95
saturation_percentile (`int`, *optional*, defaults to `95`):
Specifies the percentile uncertainty value visualized with maximum intensity.
Returns: # 返回值说明
`PIL.Image.Image` or `List[PIL.Image.Image]` with uncertainty visualization.
"""
# 定义内部函数，用于可视化单张不确定性图
def visualize_uncertainty_one(img, idx=None):
# 构建图像前缀，包含索引（如果提供）
prefix = "Uncertainty" + (f"[{idx}]" if idx else "")
# 检查图像最小值是否小于 0，若是则抛出异常
if img.min() < 0:
raise ValueError(f"{prefix}: unexected data range, min={img.min()}.")
# 将图像张量降维并转换为 NumPy 数组
img = img.squeeze(0).cpu().numpy()
# 计算图像的饱和度值，基于给定的百分位数
saturation_value = np.percentile(img, saturation_percentile)
# 将图像值归一化并限制在 0 到 255 之间
img = np.clip(img * 255 / saturation_value, 0, 255)
# 将图像数据类型转换为无符号整型（uint8）
img = img.astype(np.uint8)
# 从 NumPy 数组创建 PIL 图像对象
img = PIL.Image.fromarray(img)
# 返回处理后的图像
return img
# 检查不确定性输入是否为 None 或者是包含 None 的列表
if uncertainty is None or isinstance(uncertainty, list) and any(o is None for o in uncertainty):
# 抛出异常，输入不确定性为 None
raise ValueError("Input uncertainty is `None`")
# 如果不确定性是 NumPy 数组或 PyTorch 张量
if isinstance(uncertainty, (np.ndarray, torch.Tensor)):
# 扩展张量或数组以适应处理
uncertainty = MarigoldImageProcessor.expand_tensor_or_array(uncertainty)
# 如果不确定性为 NumPy 数组，将其转换为 PyTorch 张量
if isinstance(uncertainty, np.ndarray):
uncertainty = MarigoldImageProcessor.numpy_to_pt(uncertainty) # [N,1,H,W]
# 检查不确定性数组的维度和形状是否符合预期
if not (uncertainty.ndim == 4 and uncertainty.shape[1] == 1):
# 抛出异常，形状不符合预期
raise ValueError(f"Unexpected input shape={uncertainty.shape}, expecting [N,1,H,W].")
# 返回每个图像的可视化结果，生成图像列表
return [visualize_uncertainty_one(img, idx) for idx, img in enumerate(uncertainty)]
# 如果不确定性是一个列表
elif isinstance(uncertainty, list):
# 返回每个图像的可视化结果，生成图像列表
return [visualize_uncertainty_one(img, idx) for idx, img in enumerate(uncertainty)]
else:
# 抛出异常，输入类型不符合预期
raise ValueError(f"Unexpected input type: {type(uncertainty)}")
.\diffusers\pipelines\marigold\pipeline_marigold_depth.py
# 版权声明，说明该代码的版权归属
# Copyright 2024 Marigold authors, PRS ETH Zurich. All rights reserved.
# 版权声明，说明该代码的版权归属
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 根据 Apache 2.0 许可证对该文件的使用进行说明
# Licensed under the Apache License, Version 2.0 (the "License");
# 说明除非遵循许可证，否则不可使用该文件
# you may not use this file except in compliance with the License.
# 提供获取许可证的链接
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 提供关于该软件的使用条款
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 不提供任何类型的保证或条件
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 提供许可证的详细信息
# See the License for the specific language governing permissions and
# limitations under the License.
# --------------------------------------------------------------------------
# 提供更多信息和引用说明的来源
# More information and citation instructions are available on the
# Marigold project website: https://marigoldmonodepth.github.io
# --------------------------------------------------------------------------
# 从 dataclass 模块导入 dataclass 装饰器
from dataclasses import dataclass
# 从 functools 模块导入 partial 函数
from functools import partial
# 导入用于类型提示的类型
from typing import Any, Dict, List, Optional, Tuple, Union
# 导入 numpy 库
import numpy as np
# 导入 torch 库
import torch
# 从 PIL 导入 Image 类
from PIL import Image
# 从 tqdm 导入进度条显示工具
from tqdm.auto import tqdm
# 从 transformers 导入 CLIP 模型和标记器
from transformers import CLIPTextModel, CLIPTokenizer
# 从图像处理模块导入 PipelineImageInput 类
from ...image_processor import PipelineImageInput
# 从模型模块导入相关模型
from ...models import (
AutoencoderKL,
UNet2DConditionModel,
)
# 从调度器模块导入调度器类
from ...schedulers import (
DDIMScheduler,
LCMScheduler,
)
# 从工具模块导入基本输出和日志功能
from ...utils import (
BaseOutput,
logging,
replace_example_docstring,
)
# 导入 SciPy 可用性检查函数
from ...utils.import_utils import is_scipy_available
# 导入生成随机张量的工具
from ...utils.torch_utils import randn_tensor
# 从管道工具模块导入 DiffusionPipeline 类
from ..pipeline_utils import DiffusionPipeline
# 从图像处理模块导入 MarigoldImageProcessor 类
from .marigold_image_processing import MarigoldImageProcessor
# 创建一个日志记录器实例
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，展示如何使用该管道
EXAMPLE_DOC_STRING = """
Examples:
>>> import diffusers
>>> import torch
>>> pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
... "prs-eth/marigold-depth-lcm-v1-0", variant="fp16", torch_dtype=torch.float16
... ).to("cuda")
>>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
>>> depth = pipe(image)
>>> vis = pipe.image_processor.visualize_depth(depth.prediction)
>>> vis[0].save("einstein_depth.png")
>>> depth_16bit = pipe.image_processor.export_depth_to_16bit_png(depth.prediction)
>>> depth_16bit[0].save("einstein_depth_16bit.png")
"""
# 定义 MarigoldDepthOutput 类，表示单目深度预测的输出
@dataclass
class MarigoldDepthOutput(BaseOutput):
"""
Output class for Marigold monocular depth prediction pipeline.
# 函数参数文档字符串，说明参数的类型和形状
Args:
prediction (`np.ndarray`, `torch.Tensor`): # 预测的深度图，值范围在 [0, 1] 之间
# 形状为 numimages × 1 × height × width，无论图像是作为 4D 数组还是列表传入
uncertainty (`None`, `np.ndarray`, `torch.Tensor`): # 从集成计算的置信度图，值范围在 [0, 1] 之间
# 形状为 numimages × 1 × height × width
latent (`None`, `torch.Tensor`): # 与预测对应的潜在特征，兼容于管道的 latents 参数
# 形状为 numimages * numensemble × 4 × latentheight × latentwidth
""" # 结束文档字符串
prediction: Union[np.ndarray, torch.Tensor] # 声明 prediction 为 np.ndarray 或 torch.Tensor 类型
uncertainty: Union[None, np.ndarray, torch.Tensor] # 声明 uncertainty 为 None、np.ndarray 或 torch.Tensor 类型
latent: Union[None, torch.Tensor] # 声明 latent 为 None 或 torch.Tensor 类型
# 定义一个名为 MarigoldDepthPipeline 的类，继承自 DiffusionPipeline
class MarigoldDepthPipeline(DiffusionPipeline):
"""
使用 Marigold 方法进行单目深度估计的管道： https://marigoldmonodepth.github.io。
此模型继承自 [`DiffusionPipeline`]。请查阅父类文档，以了解库为所有管道实现的通用方法
（例如下载或保存，在特定设备上运行等）。
参数：
unet (`UNet2DConditionModel`):
条件 U-Net，用于在图像潜在空间的条件下去噪深度潜在。
vae (`AutoencoderKL`):
变分自编码器（VAE）模型，用于将图像和预测编码和解码为潜在表示。
scheduler (`DDIMScheduler` 或 `LCMScheduler`):
用于与 `unet` 结合使用的调度器，以去噪编码的图像潜在。
text_encoder (`CLIPTextModel`):
文本编码器，用于空文本嵌入。
tokenizer (`CLIPTokenizer`):
CLIP 分词器。
prediction_type (`str`, *可选*):
模型所做预测的类型。
scale_invariant (`bool`, *可选*):
指定预测深度图是否具有尺度不变性的模型属性。此值必须在模型配置中设置。
与 `shift_invariant=True` 标志一起使用时，模型也被称为“仿射不变”。注意：不支持覆盖此值。
shift_invariant (`bool`, *可选*):
指定预测深度图是否具有平移不变性的模型属性。此值必须在模型配置中设置。
与 `scale_invariant=True` 标志一起使用时，模型也被称为“仿射不变”。注意：不支持覆盖此值。
default_denoising_steps (`int`, *可选*):
生成合理质量预测所需的最小去噪扩散步骤数。此值必须在模型配置中设置。
当调用管道而未显式设置 `num_inference_steps` 时，使用默认值。这是为了确保与各种模型
变体兼容的合理结果，例如依赖于非常短的去噪调度的模型（`LCMScheduler`）和具有完整扩散
调度的模型（`DDIMScheduler`）。
default_processing_resolution (`int`, *可选*):
管道 `processing_resolution` 参数的推荐值。此值必须在模型配置中设置。
当调用管道而未显式设置 `processing_resolution` 时，使用默认值。这是为了确保与训练
使用不同最佳处理分辨率值的各种模型变体兼容的合理结果。
"""
# 定义模型在 CPU 上的卸载顺序
model_cpu_offload_seq = "text_encoder->unet->vae"
# 定义支持的预测类型，包括深度和视差
supported_prediction_types = ("depth", "disparity")
# 初始化方法，设置模型的基本参数
def __init__(
# UNet2DConditionModel 对象，进行图像生成
self,
unet: UNet2DConditionModel,
# 自动编码器，用于处理图像
vae: AutoencoderKL,
# 调度器，控制生成过程中的步伐
scheduler: Union[DDIMScheduler, LCMScheduler],
# 文本编码器，处理输入文本信息
text_encoder: CLIPTextModel,
# 分词器，用于将文本转换为模型可处理的格式
tokenizer: CLIPTokenizer,
# 可选的预测类型，默认为 None
prediction_type: Optional[str] = None,
# 可选，指示是否使用尺度不变性，默认为 True
scale_invariant: Optional[bool] = True,
# 可选，指示是否使用平移不变性，默认为 True
shift_invariant: Optional[bool] = True,
# 可选，默认去噪步骤数，默认为 None
default_denoising_steps: Optional[int] = None,
# 可选，默认处理分辨率，默认为 None
default_processing_resolution: Optional[int] = None,
):
# 调用父类的初始化方法
super().__init__()
# 检查给定的预测类型是否在支持的范围内
if prediction_type not in self.supported_prediction_types:
# 记录警告，提示可能使用了不支持的预测类型
logger.warning(
f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: "
f"{self.supported_prediction_types}."
)
# 注册模型组件，包括 UNet、VAE、调度器、文本编码器和分词器
self.register_modules(
unet=unet,
vae=vae,
scheduler=scheduler,
text_encoder=text_encoder,
tokenizer=tokenizer,
)
# 将配置参数注册到当前实例
self.register_to_config(
prediction_type=prediction_type,
scale_invariant=scale_invariant,
shift_invariant=shift_invariant,
default_denoising_steps=default_denoising_steps,
default_processing_resolution=default_processing_resolution,
)
# 计算 VAE 的缩放因子，基于其配置
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
# 设置尺度不变性
self.scale_invariant = scale_invariant
# 设置平移不变性
self.shift_invariant = shift_invariant
# 设置默认去噪步骤数
self.default_denoising_steps = default_denoising_steps
# 设置默认处理分辨率
self.default_processing_resolution = default_processing_resolution
# 初始化空文本嵌入，初始值为 None
self.empty_text_embedding = None
# 创建图像处理器，基于 VAE 的缩放因子
self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)
# 检查输入的有效性的方法
def check_inputs(
# 输入的图像，类型为管道图像输入
self,
image: PipelineImageInput,
# 推理步骤的数量
num_inference_steps: int,
# 集成的大小
ensemble_size: int,
# 处理分辨率
processing_resolution: int,
# 输入的重采样方法
resample_method_input: str,
# 输出的重采样方法
resample_method_output: str,
# 批量大小
batch_size: int,
# 可选的集成参数
ensembling_kwargs: Optional[Dict[str, Any]],
# 可选的潜在变量
latents: Optional[torch.Tensor],
# 可选的随机数生成器，支持单个或列表形式
generator: Optional[Union[torch.Generator, List[torch.Generator]]],
# 输出类型的字符串
output_type: str,
# 是否输出不确定性，布尔值
output_uncertainty: bool,
# 定义一个进度条的方法，接受可选参数用于控制进度条的显示
def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
# 检查实例是否已有进度条配置属性
if not hasattr(self, "_progress_bar_config"):
# 如果没有，初始化一个空字典作为配置
self._progress_bar_config = {}
# 如果已有配置，检查其是否为字典类型
elif not isinstance(self._progress_bar_config, dict):
# 如果不是，抛出类型错误
raise ValueError(
f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
)
# 创建一个进度条配置字典，复制已有配置
progress_bar_config = dict(**self._progress_bar_config)
# 从配置中获取描述，如果没有则使用传入的描述
progress_bar_config["desc"] = progress_bar_config.get("desc", desc)
# 从配置中获取是否保留进度条，默认值为传入的参数
progress_bar_config["leave"] = progress_bar_config.get("leave", leave)
# 如果提供了可迭代对象，返回带进度条的迭代器
if iterable is not None:
return tqdm(iterable, **progress_bar_config)
# 如果提供了总数，返回总数进度条
elif total is not None:
return tqdm(total=total, **progress_bar_config)
# 如果两个参数都未提供，抛出错误
else:
raise ValueError("Either `total` or `iterable` has to be defined.")
# 禁用梯度计算以节省内存
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义调用方法，接受多个参数以处理图像
def __call__(
self,
image: PipelineImageInput, # 输入图像
num_inference_steps: Optional[int] = None, # 推理步骤的数量
ensemble_size: int = 1, # 集成模型的数量
processing_resolution: Optional[int] = None, # 处理图像的分辨率
match_input_resolution: bool = True, # 是否匹配输入分辨率
resample_method_input: str = "bilinear", # 输入图像的重采样方法
resample_method_output: str = "bilinear", # 输出图像的重采样方法
batch_size: int = 1, # 批处理的大小
ensembling_kwargs: Optional[Dict[str, Any]] = None, # 集成模型的额外参数
latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, # 潜在变量
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, # 随机数生成器
output_type: str = "np", # 输出类型，默认是 NumPy
output_uncertainty: bool = False, # 是否输出不确定性
output_latent: bool = False, # 是否输出潜在变量
return_dict: bool = True, # 是否以字典形式返回结果
# 定义准备潜在变量的方法
def prepare_latents(
self,
image: torch.Tensor, # 输入图像的张量表示
latents: Optional[torch.Tensor], # 潜在变量的张量
generator: Optional[torch.Generator], # 随机数生成器
ensemble_size: int, # 集成模型的数量
batch_size: int, # 批处理的大小
# 返回两个张量的元组
) -> Tuple[torch.Tensor, torch.Tensor]:
# 定义一个函数来提取潜在向量
def retrieve_latents(encoder_output):
# 检查 encoder_output 是否具有 latent_dist 属性
if hasattr(encoder_output, "latent_dist"):
# 返回潜在分布的众数
return encoder_output.latent_dist.mode()
# 检查 encoder_output 是否具有 latents 属性
elif hasattr(encoder_output, "latents"):
# 返回潜在向量
return encoder_output.latents
# 如果没有找到潜在向量，则抛出异常
else:
raise AttributeError("Could not access latents of provided encoder_output")
# 将编码后的图像潜在向量按批次拼接在一起
image_latent = torch.cat(
[
# 对每个批次的图像调用 retrieve_latents 函数
retrieve_latents(self.vae.encode(image[i : i + batch_size]))
for i in range(0, image.shape[0], batch_size)
],
dim=0,
) # 结果形状为 [N,4,h,w]
# 将图像潜在向量乘以缩放因子
image_latent = image_latent * self.vae.config.scaling_factor
# 在第0维重复潜在向量以适应集成大小
image_latent = image_latent.repeat_interleave(ensemble_size, dim=0) # 结果形状为 [N*E,4,h,w]
# 将潜在预测初始化为 latents
pred_latent = latents
# 如果预测潜在向量为空，生成随机张量
if pred_latent is None:
pred_latent = randn_tensor(
# 生成与图像潜在向量相同形状的随机张量
image_latent.shape,
generator=generator,
device=image_latent.device,
dtype=image_latent.dtype,
) # 结果形状为 [N*E,4,h,w]
# 返回图像潜在向量和预测潜在向量
return image_latent, pred_latent
# 解码潜在预测，返回张量
def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor:
# 检查预测潜在向量的维度和形状是否符合预期
if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels:
# 抛出值错误，如果形状不匹配
raise ValueError(
f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}."
)
# 解码预测潜在向量，返回字典中的第一个元素
prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0] # 结果形状为 [B,3,H,W]
# 计算预测的均值，保持维度
prediction = prediction.mean(dim=1, keepdim=True) # 结果形状为 [B,1,H,W]
# 将预测限制在 [-1.0, 1.0] 的范围内
prediction = torch.clip(prediction, -1.0, 1.0) # 结果形状为 [B,1,H,W]
# 将预测从 [-1, 1] 转换到 [0, 1]
prediction = (prediction + 1.0) / 2.0
# 返回最终预测结果
return prediction # 结果形状为 [B,1,H,W]
# 定义一个静态方法，用于处理深度信息
@staticmethod
def ensemble_depth(
# 输入深度张量
depth: torch.Tensor,
# 是否使用尺度不变性
scale_invariant: bool = True,
# 是否使用位移不变性
shift_invariant: bool = True,
# 是否输出不确定性
output_uncertainty: bool = False,
# 指定聚合方式，默认为中位数
reduction: str = "median",
# 正则化强度
regularizer_strength: float = 0.02,
# 最大迭代次数
max_iter: int = 2,
# 收敛容忍度
tol: float = 1e-3,
# 最大分辨率
max_res: int = 1024,
.\diffusers\pipelines\marigold\pipeline_marigold_normals.py
# 版权所有声明，说明作者和版权信息
# Copyright 2024 Marigold authors, PRS ETH Zurich. All rights reserved.
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 根据 Apache 2.0 许可证进行授权
# Licensed under the Apache License, Version 2.0 (the "License");
# 您只能在遵守许可证的情况下使用此文件
# you may not use this file except in compliance with the License.
# 您可以在此处获取许可证副本
# http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或书面同意，软件按“原样”分发，没有任何明示或暗示的担保或条件
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 不提供任何形式的保证或条件
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 查看许可证以获取有关权限和限制的具体信息
# See the License for the specific language governing permissions and
# limitations under the License.
# --------------------------------------------------------------------------
# 额外信息和引用说明可在 Marigold 项目网站上找到
# More information and citation instructions are available on the
# Marigold project website: https://marigoldmonodepth.github.io
# --------------------------------------------------------------------------
# 导入数据类装饰器
from dataclasses import dataclass
# 导入类型提示相关的类型
from typing import Any, Dict, List, Optional, Tuple, Union
# 导入 numpy 库
import numpy as np
# 导入 PyTorch 库
import torch
# 导入图像处理库 PIL
from PIL import Image
# 导入进度条库
from tqdm.auto import tqdm
# 导入 CLIP 模型和分词器
from transformers import CLIPTextModel, CLIPTokenizer
# 导入图像处理的管道输入
from ...image_processor import PipelineImageInput
# 导入自动编码器和 UNet 模型
from ...models import (
AutoencoderKL,
UNet2DConditionModel,
)
# 导入调度器
from ...schedulers import (
DDIMScheduler,
LCMScheduler,
)
# 导入工具函数
from ...utils import (
BaseOutput,
logging,
replace_example_docstring,
)
# 导入随机张量生成工具
from ...utils.torch_utils import randn_tensor
# 导入扩散管道工具
from ..pipeline_utils import DiffusionPipeline
# 导入 Marigold 图像处理工具
from .marigold_image_processing import MarigoldImageProcessor
# 创建日志记录器，用于记录日志信息
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文档字符串，提供用法示例
EXAMPLE_DOC_STRING = """
Examples:
>>> import diffusers
>>> import torch
>>> pipe = diffusers.MarigoldNormalsPipeline.from_pretrained(
... "prs-eth/marigold-normals-lcm-v0-1", variant="fp16", torch_dtype=torch.float16
... ).to("cuda")
>>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
>>> normals = pipe(image)
>>> vis = pipe.image_processor.visualize_normals(normals.prediction)
>>> vis[0].save("einstein_normals.png")
"""
# 定义 Marigold 单目法线预测管道的输出类
@dataclass
class MarigoldNormalsOutput(BaseOutput):
"""
Marigold 单目法线预测管道的输出类
# 定义函数参数的文档字符串
Args:
# 预测法线的参数，类型可以是 numpy 数组或 PyTorch 张量
prediction (`np.ndarray`, `torch.Tensor`):
# 预测的法线值范围在 [-1, 1] 之间，形状为 $numimages \times 3 \times height
\times width$，无论图像是作为 4D 数组还是列表传递。
# 不确定性地图的参数，类型可以是 None、numpy 数组或 PyTorch 张量
uncertainty (`None`, `np.ndarray`, `torch.Tensor`):
# 从集合中计算得到的不确定性地图，值范围在 [0, 1] 之间，形状为 $numimages
\times 1 \times height \times width$。
# 潜在特征的参数，类型可以是 None 或 PyTorch 张量
latent (`None`, `torch.Tensor`):
# 与预测相对应的潜在特征，兼容于管道的 `latents` 参数。
# 形状为 $numimages * numensemble \times 4 \times latentheight \times latentwidth$。
"""
# 声明预测参数的类型，支持 numpy 数组或 PyTorch 张量
prediction: Union[np.ndarray, torch.Tensor]
# 声明不确定性参数的类型，支持 None、numpy 数组或 PyTorch 张量
uncertainty: Union[None, np.ndarray, torch.Tensor]
# 声明潜在特征参数的类型，支持 None 或 PyTorch 张量
latent: Union[None, torch.Tensor]
# 定义一个名为 MarigoldNormalsPipeline 的类，继承自 DiffusionPipeline
class MarigoldNormalsPipeline(DiffusionPipeline):
"""
使用 Marigold 方法进行单目法线估计的管道： https://marigoldmonodepth.github.io.
此模型继承自 [`DiffusionPipeline`]。请查看父类文档，以了解库为所有管道实现的通用方法
（例如下载或保存、在特定设备上运行等）。
参数：
unet (`UNet2DConditionModel`):
条件 U-Net，用于在图像潜在空间条件下对法线潜在进行去噪。
vae (`AutoencoderKL`):
变分自编码器 (VAE) 模型，用于将图像和预测编码和解码为潜在表示。
scheduler (`DDIMScheduler` 或 `LCMScheduler`):
与 `unet` 结合使用的调度器，用于对编码的图像潜在进行去噪。
text_encoder (`CLIPTextModel`):
文本编码器，用于生成空的文本嵌入。
tokenizer (`CLIPTokenizer`):
CLIP 令牌化器。
prediction_type (`str`, *可选*):
模型生成的预测类型。
use_full_z_range (`bool`, *可选*):
该模型预测的法线是否使用 Z 维度的完整范围，还是仅使用其正半。
default_denoising_steps (`int`, *可选*):
生成合理质量预测所需的最小去噪扩散步骤数。此值必须在模型配置中设置。当
管道被调用而未明确设置 `num_inference_steps` 时，使用默认值。这样可以确保
与与管道兼容的各种模型版本产生合理结果，例如依赖非常短去噪调度的模型
(`LCMScheduler`) 和具有完整扩散调度的模型 (`DDIMScheduler`)。
default_processing_resolution (`int`, *可选*):
管道的 `processing_resolution` 参数的推荐值。此值必须在模型配置中设置。当
管道被调用而未明确设置 `processing_resolution` 时，使用默认值。这样可以确保
与不同最佳处理分辨率值训练的各种模型版本产生合理结果。
"""
# 定义模型的 CPU 卸载顺序
model_cpu_offload_seq = "text_encoder->unet->vae"
# 支持的预测类型，当前仅支持 "normals"
supported_prediction_types = ("normals",)
# 初始化方法，接受多个参数以配置模型
def __init__(
# 条件 U-Net 模型
self,
unet: UNet2DConditionModel,
# 变分自编码器模型
vae: AutoencoderKL,
# 调度器，可以是 DDIMScheduler 或 LCMScheduler
scheduler: Union[DDIMScheduler, LCMScheduler],
# 文本编码器
text_encoder: CLIPTextModel,
# CLIP 令牌化器
tokenizer: CLIPTokenizer,
# 可选的预测类型
prediction_type: Optional[str] = None,
# 是否使用 Z 维度的完整范围
use_full_z_range: Optional[bool] = True,
# 默认去噪步骤数
default_denoising_steps: Optional[int] = None,
# 默认处理分辨率
default_processing_resolution: Optional[int] = None,
):
# 调用父类的构造函数
super().__init__()
# 检查预测类型是否在支持的类型中
if prediction_type not in self.supported_prediction_types:
# 如果不支持，记录警告信息
logger.warning(
f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: "
f"{self.supported_prediction_types}."
)
# 注册各个模块到当前实例中
self.register_modules(
unet=unet, # 注册 UNet 模块
vae=vae, # 注册变分自编码器模块
scheduler=scheduler, # 注册调度器模块
text_encoder=text_encoder, # 注册文本编码器模块
tokenizer=tokenizer, # 注册分词器模块
)
# 将配置参数注册到当前实例中
self.register_to_config(
use_full_z_range=use_full_z_range, # 注册使用完整的 z 范围
default_denoising_steps=default_denoising_steps, # 注册默认去噪步骤
default_processing_resolution=default_processing_resolution, # 注册默认处理分辨率
)
# 计算 VAE 的缩放因子，基于块输出通道的数量
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
# 将使用完整 z 范围的设置保存到实例中
self.use_full_z_range = use_full_z_range
# 将默认去噪步骤保存到实例中
self.default_denoising_steps = default_denoising_steps
# 将默认处理分辨率保存到实例中
self.default_processing_resolution = default_processing_resolution
# 初始化空的文本嵌入
self.empty_text_embedding = None
# 创建图像处理器实例，并传入 VAE 缩放因子
self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)
def check_inputs(
# 定义检查输入参数的方法，接收多个参数
image: PipelineImageInput, # 输入图像
num_inference_steps: int, # 推理步骤数
ensemble_size: int, # 集成大小
processing_resolution: int, # 处理分辨率
resample_method_input: str, # 输入重采样方法
resample_method_output: str, # 输出重采样方法
batch_size: int, # 批处理大小
ensembling_kwargs: Optional[Dict[str, Any]], # 集成相关的关键字参数
latents: Optional[torch.Tensor], # 潜在变量
generator: Optional[Union[torch.Generator, List[torch.Generator]]], # 随机数生成器
output_type: str, # 输出类型
output_uncertainty: bool, # 是否输出不确定性
def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
# 检查是否已经初始化了进度条配置
if not hasattr(self, "_progress_bar_config"):
# 如果没有，初始化为空字典
self._progress_bar_config = {}
# 如果存在配置，但不是字典类型，抛出错误
elif not isinstance(self._progress_bar_config, dict):
raise ValueError(
f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
)
# 复制当前的进度条配置
progress_bar_config = dict(**self._progress_bar_config)
# 设置描述，如果未提供则使用现有值
progress_bar_config["desc"] = progress_bar_config.get("desc", desc)
# 设置是否在完成后保留进度条
progress_bar_config["leave"] = progress_bar_config.get("leave", leave)
# 如果提供了可迭代对象，返回带进度条的可迭代对象
if iterable is not None:
return tqdm(iterable, **progress_bar_config)
# 如果提供了总数，返回带进度条的总数
elif total is not None:
return tqdm(total=total, **progress_bar_config)
# 如果两者都未提供，抛出错误
else:
raise ValueError("Either `total` or `iterable` has to be defined.")
# 使用 torch.no_grad() 装饰器，禁用梯度计算
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义可调用的类方法，处理图像并生成推理结果
def __call__(
self,
image: PipelineImageInput, # 输入的图像数据
num_inference_steps: Optional[int] = None, # 推理步骤数，默认为 None
ensemble_size: int = 1, # 集成模型的大小，默认为 1
processing_resolution: Optional[int] = None, # 处理分辨率，默认为 None
match_input_resolution: bool = True, # 是否匹配输入分辨率
resample_method_input: str = "bilinear", # 输入重采样方法，默认为双线性
resample_method_output: str = "bilinear", # 输出重采样方法，默认为双线性
batch_size: int = 1, # 批处理大小，默认为 1
ensembling_kwargs: Optional[Dict[str, Any]] = None, # 集成参数，默认为 None
latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, # 潜在变量，默认为 None
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, # 随机数生成器，默认为 None
output_type: str = "np", # 输出类型，默认为 NumPy
output_uncertainty: bool = False, # 是否输出不确定性，默认为 False
output_latent: bool = False, # 是否输出潜在变量，默认为 False
return_dict: bool = True, # 是否返回字典格式，默认为 True
# 从 diffusers.pipelines.marigold.pipeline_marigold_depth.MarigoldDepthPipeline.prepare_latents 复制
def prepare_latents(
self,
image: torch.Tensor, # 输入图像的张量格式
latents: Optional[torch.Tensor], # 潜在变量的张量格式
generator: Optional[torch.Generator], # 随机数生成器，默认为 None
ensemble_size: int, # 集成大小
batch_size: int, # 批处理大小
) -> Tuple[torch.Tensor, torch.Tensor]: # 返回两个张量的元组
def retrieve_latents(encoder_output): # 定义内部函数，用于获取潜在变量
if hasattr(encoder_output, "latent_dist"): # 检查输出是否包含潜在分布
return encoder_output.latent_dist.mode() # 返回潜在分布的众数
elif hasattr(encoder_output, "latents"): # 检查输出是否包含潜在变量
return encoder_output.latents # 返回潜在变量
else: # 如果都没有，抛出异常
raise AttributeError("Could not access latents of provided encoder_output")
image_latent = torch.cat( # 将处理后的潜在变量进行拼接
[
retrieve_latents(self.vae.encode(image[i : i + batch_size])) # 对每个批次的图像编码并获取潜在变量
for i in range(0, image.shape[0], batch_size) # 按批处理大小遍历图像
],
dim=0, # 在第0维进行拼接
) # [N,4,h,w] # 得到的潜在变量张量的形状
image_latent = image_latent * self.vae.config.scaling_factor # 应用缩放因子调整潜在变量
image_latent = image_latent.repeat_interleave(ensemble_size, dim=0) # [N*E,4,h,w] # 重复以匹配集成大小
pred_latent = latents # 初始化预测潜在变量
if pred_latent is None: # 如果未提供潜在变量
pred_latent = randn_tensor( # 生成随机潜在变量
image_latent.shape, # 形状与 image_latent 相同
generator=generator, # 使用提供的随机数生成器
device=image_latent.device, # 使用 image_latent 的设备
dtype=image_latent.dtype, # 使用 image_latent 的数据类型
) # [N*E,4,h,w] # 生成的潜在变量形状
return image_latent, pred_latent # 返回图像潜在变量和预测潜在变量
def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor: # 解码预测潜在变量的方法
if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels: # 检查预测潜在变量的维度和通道数
raise ValueError( # 如果不符合要求，抛出异常
f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}."
)
prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0] # [B,3,H,W] # 解码潜在变量，得到预测图像
prediction = torch.clip(prediction, -1.0, 1.0) # 限制预测值在 -1.0 到 1.0 之间
if not self.use_full_z_range: # 如果不使用完整的潜在范围
prediction[:, 2, :, :] *= 0.5 # 对第三个通道进行缩放
prediction[:, 2, :, :] += 0.5 # 对第三个通道进行偏移
prediction = self.normalize_normals(prediction) # [B,3,H,W] # 正常化预测结果
return prediction # [B,3,H,W] # 返回最终的预测图像
@staticmethod # 静态方法的标记
# 规范化法线向量，使其单位长度，避免数值不稳定
def normalize_normals(normals: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
# 检查输入的法线张量是否为4维且第二维的大小为3
if normals.dim() != 4 or normals.shape[1] != 3:
# 如果不满足条件，抛出错误
raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
# 计算法线张量在第二维的范数，并保持维度
norm = torch.norm(normals, dim=1, keepdim=True)
# 将法线张量除以范数，使用clamp限制最小值以避免除以零
normals /= norm.clamp(min=eps)
# 返回规范化后的法线张量
return normals
@staticmethod
# 对法线张量进行集成处理，返回集成后的法线和可选的不确定性
def ensemble_normals(
normals: torch.Tensor, output_uncertainty: bool, reduction: str = "closest"
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
"""
对法线图进行集成，期望输入形状为 `(B, 3, H, W)`，其中 B 是
每个预测的集成成员数量，大小为 `(H x W)`。
Args:
normals (`torch.Tensor`):
输入的集成法线图。
output_uncertainty (`bool`, *可选*, 默认值为 `False`):
是否输出不确定性图。
reduction (`str`, *可选*, 默认值为 `"closest"`):
用于集成对齐预测的归约方法。接受的值为：`"closest"` 和
`"mean"`。
Returns:
返回形状为 `(1, 3, H, W)` 的对齐和集成法线图，以及可选的不确定性张量，形状为 `(1, 1, H, W)`。
"""
# 检查输入的法线张量是否为4维且第二维的大小为3
if normals.dim() != 4 or normals.shape[1] != 3:
# 如果不满足条件，抛出错误
raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
# 检查归约方法是否有效
if reduction not in ("closest", "mean"):
# 如果不合法，抛出错误
raise ValueError(f"Unrecognized reduction method: {reduction}.")
# 计算法线的均值，保持维度
mean_normals = normals.mean(dim=0, keepdim=True) # [1,3,H,W]
# 规范化均值法线
mean_normals = MarigoldNormalsPipeline.normalize_normals(mean_normals) # [1,3,H,W]
# 计算均值法线与所有法线的点积，得到相似度
sim_cos = (mean_normals * normals).sum(dim=1, keepdim=True) # [E,1,H,W]
# 限制相似度值在 -1 到 1 之间，以避免在 fp16 中出现 NaN
sim_cos = sim_cos.clamp(-1, 1) # required to avoid NaN in uncertainty with fp16
uncertainty = None
# 如果需要输出不确定性
if output_uncertainty:
# 计算相似度的反余弦，得到不确定性
uncertainty = sim_cos.arccos() # [E,1,H,W]
# 计算均值并归一化
uncertainty = uncertainty.mean(dim=0, keepdim=True) / np.pi # [1,1,H,W]
# 如果选择平均归约方法
if reduction == "mean":
# 返回均值法线和不确定性
return mean_normals, uncertainty # [1,3,H,W], [1,1,H,W]
# 找到相似度最大的索引
closest_indices = sim_cos.argmax(dim=0, keepdim=True) # [1,1,H,W]
# 将索引扩展到法线的通道数
closest_indices = closest_indices.repeat(1, 3, 1, 1) # [1,3,H,W]
# 根据索引从法线中提取相应的法线
closest_normals = torch.gather(normals, 0, closest_indices) # [1,3,H,W]
# 返回最近法线和不确定性
return closest_normals, uncertainty # [1,3,H,W], [1,1,H,W]
.\diffusers\pipelines\marigold\__init__.py
# 导入类型检查相关的模块
from typing import TYPE_CHECKING
# 从相对路径导入所需的工具函数和常量
from ...utils import (
DIFFUSERS_SLOW_IMPORT, # 慢导入标志
OptionalDependencyNotAvailable, # 可选依赖未可用的异常类
_LazyModule, # 延迟模块加载类
get_objects_from_module, # 从模块获取对象的函数
is_torch_available, # 检查 PyTorch 是否可用的函数
is_transformers_available, # 检查 Transformers 是否可用的函数
)
# 初始化一个空字典，用于存储占位对象
_dummy_objects = {}
# 初始化一个空字典，用于存储导入结构
_import_structure = {}
# 尝试检测依赖关系
try:
# 检查 Transformers 和 PyTorch 是否都可用
if not (is_transformers_available() and is_torch_available()):
# 如果不可用，则抛出异常
raise OptionalDependencyNotAvailable()
# 捕获可选依赖未可用的异常
except OptionalDependencyNotAvailable:
# 从工具模块中导入占位对象（忽略 F403 警告）
from ...utils import dummy_torch_and_transformers_objects # noqa F403
# 更新占位对象字典
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
# 如果依赖可用，定义模块的导入结构
_import_structure["marigold_image_processing"] = ["MarigoldImageProcessor"]
_import_structure["pipeline_marigold_depth"] = ["MarigoldDepthOutput", "MarigoldDepthPipeline"]
_import_structure["pipeline_marigold_normals"] = ["MarigoldNormalsOutput", "MarigoldNormalsPipeline"]
# 检查类型检查或慢导入标志
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
# 尝试检测依赖关系
try:
# 检查 Transformers 和 PyTorch 是否都可用
if not (is_transformers_available() and is_torch_available()):
# 如果不可用，则抛出异常
raise OptionalDependencyNotAvailable()
# 捕获可选依赖未可用的异常
except OptionalDependencyNotAvailable:
# 导入占位对象以避免错误
from ...utils.dummy_torch_and_transformers_objects import *
else:
# 导入实际模块中的类
from .marigold_image_processing import MarigoldImageProcessor
from .pipeline_marigold_depth import MarigoldDepthOutput, MarigoldDepthPipeline
from .pipeline_marigold_normals import MarigoldNormalsOutput, MarigoldNormalsPipeline
else:
# 导入系统模块
import sys
# 用延迟模块加载类替代当前模块
sys.modules[__name__] = _LazyModule(
__name__,
globals()["__file__"], # 当前文件的全局名称
_import_structure, # 导入结构
module_spec=__spec__, # 模块的规范
)
# 将占位对象添加到当前模块中
for name, value in _dummy_objects.items():
setattr(sys.modules[__name__], name, value)
.\diffusers\pipelines\musicldm\pipeline_musicldm.py
# 版权声明，表明该代码的版权归 HuggingFace 团队所有
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 许可声明，指明该代码遵循 Apache 许可证 2.0 版本
# Licensed under the Apache License, Version 2.0 (the "License");
# 除非遵循许可证，否则不可使用此文件
# you may not use this file except in compliance with the License.
# 提供许可证获取链接
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# 指出软件按“原样”分发，没有任何明示或暗示的保证
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 参考许可证以了解具体权限和限制
# See the License for the specific language governing permissions and
# limitations under the License.
# 导入 inspect 模块以获取有关对象的信息
import inspect
# 从 typing 模块导入类型提示相关的类型
from typing import Any, Callable, Dict, List, Optional, Union
# 导入 numpy 库以进行数值计算
import numpy as np
# 导入 PyTorch 库以进行深度学习
import torch
# 从 transformers 库导入多个模型和特征提取器
from transformers import (
ClapFeatureExtractor, # 导入 Clap 的特征提取器
ClapModel, # 导入 Clap 模型
ClapTextModelWithProjection, # 导入带投影的 Clap 文本模型
RobertaTokenizer, # 导入 Roberta 分词器
RobertaTokenizerFast, # 导入快速 Roberta 分词器
SpeechT5HifiGan, # 导入 SpeechT5 的 HiFiGan 模型
)
# 从相对路径导入模型和调度器
from ...models import AutoencoderKL, UNet2DConditionModel # 导入自编码器和条件 UNet 模型
from ...schedulers import KarrasDiffusionSchedulers # 导入 Karras 扩散调度器
from ...utils import (
is_accelerate_available, # 检查 accelerate 是否可用
is_accelerate_version, # 检查 accelerate 版本
is_librosa_available, # 检查 librosa 是否可用
logging, # 导入日志记录功能
replace_example_docstring, # 导入替换示例文档字符串的工具
)
from ...utils.torch_utils import randn_tensor # 导入生成随机张量的工具
from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline, StableDiffusionMixin # 导入音频管道输出和扩散管道类
# 如果 librosa 库可用，则导入它
if is_librosa_available():
import librosa
# 创建一个日志记录器，用于记录当前模块的信息
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 定义示例文档字符串，展示如何使用 MusicLDMPipeline
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import MusicLDMPipeline # 从 diffusers 导入音乐生成管道
>>> import torch # 导入 PyTorch 库
>>> import scipy # 导入 SciPy 库用于处理音频
>>> repo_id = "ucsd-reach/musicldm" # 定义模型仓库 ID
>>> pipe = MusicLDMPipeline.from_pretrained(repo_id, torch_dtype=torch.float16) # 从预训练模型创建管道
>>> pipe = pipe.to("cuda") # 将管道移至 CUDA 设备
>>> prompt = "Techno music with a strong, upbeat tempo and high melodic riffs" # 定义生成音乐的提示
>>> audio = pipe(prompt, num_inference_steps=10, audio_length_in_s=5.0).audios[0] # 生成音频
>>> # 将生成的音频保存为 .wav 文件
>>> scipy.io.wavfile.write("techno.wav", rate=16000, data=audio) # 保存音频文件
```py
"""
# 定义 MusicLDMPipeline 类，继承自 DiffusionPipeline 和 StableDiffusionMixin
class MusicLDMPipeline(DiffusionPipeline, StableDiffusionMixin):
r"""
用于基于文本生成音频的管道，使用 MusicLDM 模型。
该模型继承自 [`DiffusionPipeline`]。请查看超类文档以获取所有管道实现的通用方法
(下载、保存、在特定设备上运行等)。
# 文档字符串，描述构造函数的参数及其类型
Args:
vae ([`AutoencoderKL`]):
# 变分自编码器（VAE）模型，用于将图像编码和解码为潜在表示
text_encoder ([`~transformers.ClapModel`]):
# 冻结的文本-音频嵌入模型（`ClapTextModel`），特别是
# [laion/clap-htsat-unfused](https://huggingface.co/laion/clap-htsat-unfused) 变体
tokenizer ([`PreTrainedTokenizer`]):
# [`~transformers.RobertaTokenizer`] 用于对文本进行分词
feature_extractor ([`~transformers.ClapFeatureExtractor`]):
# 特征提取器，用于从音频波形计算梅尔谱图
unet ([`UNet2DConditionModel`]):
# `UNet2DConditionModel` 用于去噪编码后的音频潜在表示
scheduler ([`SchedulerMixin`]):
# 调度器，与 `unet` 结合使用以去噪编码的音频潜在表示，可以是
# [`DDIMScheduler`], [`LMSDiscreteScheduler`] 或 [`PNDMScheduler`]
vocoder ([`~transformers.SpeechT5HifiGan`]):
# `SpeechT5HifiGan` 类的声码器
"""
# 构造函数，初始化对象的属性
def __init__(
self,
vae: AutoencoderKL, # 变分自编码器模型
text_encoder: Union[ClapTextModelWithProjection, ClapModel], # 文本编码器
tokenizer: Union[RobertaTokenizer, RobertaTokenizerFast], # 分词器
feature_extractor: Optional[ClapFeatureExtractor], # 可选的特征提取器
unet: UNet2DConditionModel, # UNet 模型用于去噪
scheduler: KarrasDiffusionSchedulers, # 调度器
vocoder: SpeechT5HifiGan, # 声码器
):
super().__init__() # 调用父类的构造函数
# 注册模块，存储各种组件到对象的属性中
self.register_modules(
vae=vae, # 注册 VAE 模型
text_encoder=text_encoder, # 注册文本编码器
tokenizer=tokenizer, # 注册分词器
feature_extractor=feature_extractor, # 注册特征提取器
unet=unet, # 注册 UNet 模型
scheduler=scheduler, # 注册调度器
vocoder=vocoder, # 注册声码器
)
# 计算 VAE 的缩放因子，根据 VAE 配置的块输出通道数计算
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
# 编码提示的私有方法，处理输入提示
def _encode_prompt(
self,
prompt, # 输入提示文本
device, # 设备信息
num_waveforms_per_prompt, # 每个提示生成的波形数量
do_classifier_free_guidance, # 是否进行无分类器引导
negative_prompt=None, # 可选的负面提示
prompt_embeds: Optional[torch.Tensor] = None, # 可选的提示嵌入
negative_prompt_embeds: Optional[torch.Tensor] = None, # 可选的负面提示嵌入
# 从 diffusers.pipelines.audioldm.pipeline_audioldm.AudioLDMPipeline 复制的梅尔谱图到波形的转换方法
def mel_spectrogram_to_waveform(self, mel_spectrogram): # 定义方法，将梅尔谱图转换为波形
# 如果梅尔谱图是四维的，则去掉第一维
if mel_spectrogram.dim() == 4:
mel_spectrogram = mel_spectrogram.squeeze(1)
# 使用声码器将梅尔谱图转换为波形
waveform = self.vocoder(mel_spectrogram)
# 始终转换为 float32 类型，因为这不会造成显著的开销，并且与 bfloat16 兼容
waveform = waveform.cpu().float() # 将波形移动到 CPU 并转换为 float32 类型
return waveform # 返回生成的波形
# 从 diffusers.pipelines.audioldm2.pipeline_audioldm2.AudioLDM2Pipeline 复制的得分波形的方法
# 评分音频波形与文本提示之间的匹配度
def score_waveforms(self, text, audio, num_waveforms_per_prompt, device, dtype):
# 检查是否安装了 librosa 包
if not is_librosa_available():
# 记录信息，提示用户安装 librosa 包以启用自动评分
logger.info(
"Automatic scoring of the generated audio waveforms against the input prompt text requires the "
"`librosa` package to resample the generated waveforms. Returning the audios in the order they were "
"generated. To enable automatic scoring, install `librosa` with: `pip install librosa`."
)
# 返回原始音频
return audio
# 对文本进行标记化并返回张量格式的输入
inputs = self.tokenizer(text, return_tensors="pt", padding=True)
# 使用 librosa 对音频进行重采样，调整采样率
resampled_audio = librosa.resample(
audio.numpy(), orig_sr=self.vocoder.config.sampling_rate, target_sr=self.feature_extractor.sampling_rate
)
# 将重采样后的音频特征提取并转换为指定数据类型
inputs["input_features"] = self.feature_extractor(
list(resampled_audio), return_tensors="pt", sampling_rate=self.feature_extractor.sampling_rate
).input_features.type(dtype)
# 将输入数据移动到指定设备上
inputs = inputs.to(device)
# 使用 CLAP 模型计算音频与文本的相似性得分
logits_per_text = self.text_encoder(**inputs).logits_per_text
# 根据文本匹配度对生成的音频进行排序
indices = torch.argsort(logits_per_text, dim=1, descending=True)[:, :num_waveforms_per_prompt]
# 选择根据排序结果的音频
audio = torch.index_select(audio, 0, indices.reshape(-1).cpu())
# 返回排序后的音频
return audio
# 从 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs 复制
def prepare_extra_step_kwargs(self, generator, eta):
# 准备调度器步骤的额外参数，因为并非所有调度器都有相同的参数签名
# eta (η) 仅在 DDIMScheduler 中使用，其他调度器会被忽略。
# eta 对应 DDIM 论文中的 η: https://arxiv.org/abs/2010.02502
# 其值应在 [0, 1] 之间
# 检查调度器是否接受 eta 参数
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
extra_step_kwargs = {}
# 如果接受 eta，则将其添加到额外参数中
if accepts_eta:
extra_step_kwargs["eta"] = eta
# 检查调度器是否接受 generator 参数
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 如果接受 generator，则将其添加到额外参数中
if accepts_generator:
extra_step_kwargs["generator"] = generator
# 返回包含额外参数的字典
return extra_step_kwargs
# 从 diffusers.pipelines.audioldm.pipeline_audioldm.AudioLDMPipeline.check_inputs 复制
def check_inputs(
self,
prompt,
audio_length_in_s,
vocoder_upsample_factor,
callback_steps,
negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
):
# 计算最小音频长度（秒），基于声码器上采样因子和 VAE 缩放因子
min_audio_length_in_s = vocoder_upsample_factor * self.vae_scale_factor
# 如果输入音频长度小于最小音频长度，则抛出错误
if audio_length_in_s < min_audio_length_in_s:
raise ValueError(
# 提示音频长度必须大于等于最小音频长度
f"`audio_length_in_s` has to be a positive value greater than or equal to {min_audio_length_in_s}, but "
f"is {audio_length_in_s}."
)
# 检查声码器模型输入维度是否可以被 VAE 缩放因子整除
if self.vocoder.config.model_in_dim % self.vae_scale_factor != 0:
raise ValueError(
# 提示频率 bins 数量必须可以被 VAE 缩放因子整除
f"The number of frequency bins in the vocoder's log-mel spectrogram has to be divisible by the "
f"VAE scale factor, but got {self.vocoder.config.model_in_dim} bins and a scale factor of "
f"{self.vae_scale_factor}."
)
# 检查 callback_steps 是否有效（必须为正整数）
if (callback_steps is None) or (
callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
):
raise ValueError(
# 提示 callback_steps 必须为正整数
f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
f" {type(callback_steps)}."
)
# 检查同时传入 prompt 和 prompt_embeds 是否有效
if prompt is not None and prompt_embeds is not None:
raise ValueError(
# 提示不能同时提供 prompt 和 prompt_embeds
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
# 检查 prompt 和 prompt_embeds 是否同时为 None
elif prompt is None and prompt_embeds is None:
raise ValueError(
# 提示必须提供 prompt 或 prompt_embeds，不能都为空
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
# 检查 prompt 的类型是否为字符串或列表
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
# 检查是否同时传入 negative_prompt 和 negative_prompt_embeds
if negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
# 提示不能同时提供 negative_prompt 和 negative_prompt_embeds
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
# 检查 prompt_embeds 和 negative_prompt_embeds 是否形状一致
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
# 提示 prompt_embeds 和 negative_prompt_embeds 形状必须相同
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
# 从 diffusers.pipelines.audioldm.pipeline_audioldm.AudioLDMPipeline.prepare_latents 中复制的代码
# 准备潜在变量，参数包括批次大小、通道数、高度等
def prepare_latents(self, batch_size, num_channels_latents, height, dtype, device, generator, latents=None):
# 定义潜在变量的形状，考虑 VAE 的缩放因子
shape = (
batch_size,
num_channels_latents,
int(height) // self.vae_scale_factor,
int(self.vocoder.config.model_in_dim) // self.vae_scale_factor,
)
# 检查生成器列表的长度是否与批次大小匹配
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
# 如果没有提供潜在变量，则随机生成
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
# 如果提供了潜在变量，将其移动到指定设备
latents = latents.to(device)
# 根据调度器需要的标准差缩放初始噪声
latents = latents * self.scheduler.init_noise_sigma
# 返回处理后的潜在变量
return latents
# 启用模型的 CPU 卸载，以减少内存使用
def enable_model_cpu_offload(self, gpu_id=0):
r"""
Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
`enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
"""
# 检查 accelerate 是否可用且版本是否合适
if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
from accelerate import cpu_offload_with_hook
else:
# 抛出错误以提示用户需要更新 accelerate
raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")
# 设置设备为指定的 GPU
device = torch.device(f"cuda:{gpu_id}")
# 如果当前设备不是 CPU，则将模型移动到 CPU
if self.device.type != "cpu":
self.to("cpu", silence_dtype_warnings=True)
# 清空 GPU 缓存以查看内存节省
torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
# 定义需要卸载到 CPU 的模型序列
model_sequence = [
self.text_encoder.text_model,
self.text_encoder.text_projection,
self.unet,
self.vae,
self.vocoder,
self.text_encoder,
]
hook = None
# 遍历模型序列，逐个卸载到 CPU
for cpu_offloaded_model in model_sequence:
_, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
# 手动卸载最后一个模型
self.final_offload_hook = hook
# 禁用梯度计算，优化性能
@torch.no_grad()
# 替换示例文档字符串
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定义可调用对象的方法，允许使用不同参数进行推断
def __call__(
# 提示内容，可以是单个字符串或字符串列表
self,
prompt: Union[str, List[str]] = None,
# 音频长度，单位为秒，默认为 None 表示不限制
audio_length_in_s: Optional[float] = None,
# 推理步骤数量，默认为 200
num_inference_steps: int = 200,
# 引导比例，用于控制生成结果的引导强度，默认为 2.0
guidance_scale: float = 2.0,
# 负提示，可以是单个字符串或字符串列表，默认为 None
negative_prompt: Optional[Union[str, List[str]]] = None,
# 每个提示生成的波形数量，默认为 1
num_waveforms_per_prompt: Optional[int] = 1,
# 采样的 eta 值，默认为 0.0
eta: float = 0.0,
# 随机数生成器，可以是单个或多个 PyTorch 生成器，默认为 None
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 可选的潜在表示，默认为 None
latents: Optional[torch.Tensor] = None,
# 可选的提示嵌入，默认为 None
prompt_embeds: Optional[torch.Tensor] = None,
# 可选的负提示嵌入，默认为 None
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 是否返回字典形式的结果，默认为 True
return_dict: bool = True,
# 可选的回调函数，用于在推理过程中执行
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
# 回调执行的步骤间隔，默认为 1
callback_steps: Optional[int] = 1,
# 可选的交叉注意力参数，默认为 None
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
# 输出类型，默认为 "np"，表示返回 NumPy 数组
output_type: Optional[str] = "np",
