Marigold 用于计算机视觉任务的管道
Marigold 是一种新颖的基于扩散的密集预测方法,以及针对各种计算机视觉任务的一组管道,例如单目深度估计。
本指南将向您展示如何使用 Marigold 获取图像和视频的快速高质量预测。
每个管道都支持一项计算机视觉任务,该任务将输入 RGB 图像作为输入并生成所关注模态的预测,例如输入图像的深度图。目前,已实现以下任务
管道 | 预测的模态 | 演示 |
---|---|---|
MarigoldDepthPipeline | 深度,视差 | 快速演示 (LCM),缓慢的原始演示 (DDIM) |
MarigoldNormalsPipeline | 表面法线 | 快速演示 (LCM) |
原始检查点可以在 PRS-ETH Hugging Face 组织下找到。这些检查点旨在与 diffusers 管道和 原始代码库 一起使用。原始代码还可以用于训练新的检查点。
检查点 | 模态 | 评论 |
---|---|---|
prs-eth/marigold-v1-0 | 深度 | 第一个 Marigold Depth 检查点,预测仿射不变深度图。该检查点在基准测试中的性能在原始 论文 中进行了研究。专为在推理时与 DDIMScheduler 一起使用而设计,它需要至少 10 步才能获得可靠的预测。仿射不变深度预测在每个像素中具有 0(近平面)到 1(远平面)范围内的值;这两个平面都由模型作为推理过程的一部分进行选择。有关可视化实用程序,请参阅 MarigoldImageProcessor 参考。 |
prs-eth/marigold-depth-lcm-v1-0 | 深度 | 快速 Marigold Depth 检查点,从 prs-eth/marigold-v1-0 微调而来。专为在推理时与 LCMScheduler 一起使用而设计,它只需要 1 步就能获得可靠的预测。预测可靠性在 4 步时饱和,之后会下降。 |
prs-eth/marigold-normals-v0-1 | 法线 | Marigold Normals 管道的预览检查点。专为在推理时与 DDIMScheduler 一起使用而设计,它需要至少 10 步才能获得可靠的预测。表面法线预测是值在 -1 到 1 范围内的单位长度 3D 向量。该检查点将在 v1-0 版本发布后逐步淘汰。 |
prs-eth/marigold-normals-lcm-v0-1 | 法线 | 快速 Marigold Normals 检查点,从 prs-eth/marigold-normals-v0-1 微调而来。专为在推理时与 LCMScheduler 一起使用而设计,它只需要 1 步就能获得可靠的预测。预测可靠性在 4 步时饱和,之后会下降。该检查点将在 v1-0 版本发布后逐步淘汰。 |
以下示例主要针对深度预测,但可以普遍应用于其他支持的模态。我们展示了使用 Midjourney 生成的爱因斯坦的同一输入图像进行的预测。这使得比较各种模态和检查点之间预测的可视化变得更容易。
深度预测快速入门
要获得第一个深度预测,将 prs-eth/marigold-depth-lcm-v1-0
检查点加载到 MarigoldDepthPipeline
管道中,将图像通过管道,然后保存预测
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")
深度的可视化函数 visualize_depth()
应用 matplotlib 的颜色映射 之一(默认情况下为 Spectral
)将预测的像素值从单通道 [0, 1]
深度范围映射到 RGB 图像。使用 Spectral
颜色映射,近深度的像素被涂成红色,远像素被分配蓝色。16 位 PNG 文件存储从 [0, 1]
范围线性映射到 [0, 65535]
的单通道值。以下是原始预测和可视化预测;可以看到,黑暗区域(胡须)在可视化中更容易区分
表面法线预测快速入门
将 prs-eth/marigold-normals-lcm-v0-1
检查点加载到 MarigoldNormalsPipeline
管道中,将图像通过管道,然后保存预测
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")
表面法线的可视化函数 visualize_normals()
将像素值在 [-1, 1]
范围内的三维预测映射到 RGB 图像。可视化函数支持翻转表面法线轴,以使可视化与其他参考系选择兼容。从概念上讲,每个像素根据参考系中的表面法线向量进行绘制,其中 X
轴指向右侧,Y
轴指向向上,Z
轴指向观察者。以下是可视化的预测
在此示例中,鼻尖几乎肯定在表面上有一个点,其中表面法线向量直接指向观察者,这意味着它的坐标是 [0, 0, 1]
。该向量映射到 RGB [128, 128, 255]
,对应于紫罗兰色。类似地,图像右侧脸颊上的表面法线具有较大的 X
分量,从而增加了红色色调。指向上的肩膀上的点具有较大的 Y
,从而增加了绿色色调。
加快推理
以上快速入门代码段已经针对速度进行了优化:它们加载 LCM 检查点,使用 fp16
版本的权重和计算,并且仅执行一步去噪扩散步骤。pipe(image)
调用在 RTX 3090 GPU 上以 280 毫秒完成。在内部,输入图像使用 Stable Diffusion VAE 编码器进行编码,然后 U-Net 执行一步去噪步骤,最后,预测潜在向量使用 VAE 解码器解码为像素空间。在这种情况下,三个模块调用中有两个专用于在 LDM 的像素空间和潜在空间之间转换。由于 Marigold 的潜在空间与基本 Stable Diffusion 兼容,因此可以通过使用 SD VAE 的轻量级替代品 将管道调用速度提高 3 倍以上(在 RTX 3090 上为 85 毫秒)
import diffusers
import torch
pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
"prs-eth/marigold-depth-lcm-v1-0", variant="fp16", torch_dtype=torch.float16
).to("cuda")
+ pipe.vae = diffusers.AutoencoderTiny.from_pretrained(
+ "madebyollin/taesd", torch_dtype=torch.float16
+ ).cuda()
image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
depth = pipe(image)
如 优化 中所述,添加 torch.compile
可能会根据目标硬件进一步提高性能
import diffusers
import torch
pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
"prs-eth/marigold-depth-lcm-v1-0", variant="fp16", torch_dtype=torch.float16
).to("cuda")
+ pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
depth = pipe(image)
与深度模型的定性比较
通过上述速度优化,Marigold 可以比使用最大检查点 Depth Anything LiheYoung/depth-anything-large-hf 更快地提供更详细的预测。
最大化精度和集成
Marigold 管道内置集成机制,组合来自不同随机潜变量的多个预测。这是提高预测精度的一种蛮力方法,利用了扩散的生成特性。当 `ensemble_size` 参数设置为大于 `1` 时,集成路径会自动激活。当追求最大精度时,调整 `num_inference_steps` 和 `ensemble_size` 同时调整是有意义的。推荐值因检查点而异,但主要取决于调度器类型。集成的效果在表面法线上表现得尤为明显。
import diffusers
model_path = "prs-eth/marigold-normals-v1-0"
model_paper_kwargs = {
diffusers.schedulers.DDIMScheduler: {
"num_inference_steps": 10,
"ensemble_size": 10,
},
diffusers.schedulers.LCMScheduler: {
"num_inference_steps": 4,
"ensemble_size": 5,
},
}
image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
pipe = diffusers.MarigoldNormalsPipeline.from_pretrained(model_path).to("cuda")
pipe_kwargs = model_paper_kwargs[type(pipe.scheduler)]
depth = pipe(image, **pipe_kwargs)
vis = pipe.image_processor.visualize_normals(depth.prediction)
vis[0].save("einstein_normals.png")
可以看出,所有具有细粒度结构的区域,如头发,都获得了更保守且平均更准确的预测。这样的结果更适合对精度敏感的下游任务,例如 3D 重建。
定量评估
为了在标准排行榜和基准(如 NYU、KITTI 和其他数据集)中对 Marigold 进行定量评估,请遵循论文中概述的评估协议:加载完整的精度 fp32 模型,并对 `num_inference_steps` 和 `ensemble_size` 使用适当的值。可选地,对随机数进行种子以确保可重复性。最大化 `batch_size` 将实现最大设备利用率。
import diffusers
import torch
device = "cuda"
seed = 2024
model_path = "prs-eth/marigold-v1-0"
model_paper_kwargs = {
diffusers.schedulers.DDIMScheduler: {
"num_inference_steps": 50,
"ensemble_size": 10,
},
diffusers.schedulers.LCMScheduler: {
"num_inference_steps": 4,
"ensemble_size": 10,
},
}
image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
generator = torch.Generator(device=device).manual_seed(seed)
pipe = diffusers.MarigoldDepthPipeline.from_pretrained(model_path).to(device)
pipe_kwargs = model_paper_kwargs[type(pipe.scheduler)]
depth = pipe(image, generator=generator, **pipe_kwargs)
# evaluate metrics
使用预测不确定性
Marigold 管道内置的集成机制组合了从不同随机潜变量获得的多个预测。作为副作用,它可用于量化认知(模型)不确定性;只需指定 `ensemble_size` 大于 1 并设置 `output_uncertainty=True` 即可。结果不确定性将在输出的 `uncertainty` 字段中提供。可以可视化如下:
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,
ensemble_size=10, # any number greater than 1; higher values yield higher precision
output_uncertainty=True,
)
uncertainty = pipe.image_processor.visualize_uncertainty(depth.uncertainty)
uncertainty[0].save("einstein_depth_uncertainty.png")
不确定性的解释很简单:较高的值(白色)对应于模型难以做出一致预测的像素。显然,深度模型在具有不连续性的边缘周围最不确定,物体深度在这些边缘急剧变化。表面法线模型在细粒度结构(如头发)和黑暗区域(如衣领)中最不确定。
逐帧视频处理与时间一致性
由于 Marigold 的生成特性,每个预测都是唯一的,并由用于潜变量初始化的随机噪声定义。与传统的端到端密集回归网络相比,这成为了一个明显的缺点,如下面的视频所示:
为了解决这个问题,可以将 `latents` 参数传递给管道,该参数定义扩散的起点。根据经验,我们发现对同一个起点噪声潜变量和对应于前一帧预测的潜变量进行凸组合可以获得足够平滑的结果,如下面的代码片段所示:
import imageio
from PIL import Image
from tqdm import tqdm
import diffusers
import torch
device = "cuda"
path_in = "obama.mp4"
path_out = "obama_depth.gif"
pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
"prs-eth/marigold-depth-lcm-v1-0", variant="fp16", torch_dtype=torch.float16
).to(device)
pipe.vae = diffusers.AutoencoderTiny.from_pretrained(
"madebyollin/taesd", torch_dtype=torch.float16
).to(device)
pipe.set_progress_bar_config(disable=True)
with imageio.get_reader(path_in) as reader:
size = reader.get_meta_data()['size']
last_frame_latent = None
latent_common = torch.randn(
(1, 4, 768 * size[1] // (8 * max(size)), 768 * size[0] // (8 * max(size)))
).to(device=device, dtype=torch.float16)
out = []
for frame_id, frame in tqdm(enumerate(reader), desc="Processing Video"):
frame = Image.fromarray(frame)
latents = latent_common
if last_frame_latent is not None:
latents = 0.9 * latents + 0.1 * last_frame_latent
depth = pipe(
frame, match_input_resolution=False, latents=latents, output_latent=True
)
last_frame_latent = depth.latent
out.append(pipe.image_processor.visualize_depth(depth.prediction)[0])
diffusers.utils.export_to_gif(out, path_out, fps=reader.get_meta_data()['fps'])
在这里,扩散过程从给定的计算出的潜变量开始。管道设置 `output_latent=True` 以访问 `out.latent` 并计算它对下一帧潜变量初始化的贡献。现在结果稳定得多:
Marigold 用于 ControlNet
使用扩散模型进行深度预测的非常常见的应用是与 ControlNet 结合使用。深度清晰度在从 ControlNet 获得高质量结果方面起着至关重要的作用。如上所述与其他方法的比较所示,Marigold 在该任务中表现出色。下面的代码片段演示了如何加载图像、计算深度并将深度以兼容的格式传递给 ControlNet。
import torch
import diffusers
device = "cuda"
generator = torch.Generator(device=device).manual_seed(2024)
image = diffusers.utils.load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_depth_source.png"
)
pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
"prs-eth/marigold-depth-lcm-v1-0", torch_dtype=torch.float16, variant="fp16"
).to(device)
depth_image = pipe(image, generator=generator).prediction
depth_image = pipe.image_processor.visualize_depth(depth_image, color_map="binary")
depth_image[0].save("motorcycle_controlnet_depth.png")
controlnet = diffusers.ControlNetModel.from_pretrained(
"diffusers/controlnet-depth-sdxl-1.0", torch_dtype=torch.float16, variant="fp16"
).to(device)
pipe = diffusers.StableDiffusionXLControlNetPipeline.from_pretrained(
"SG161222/RealVisXL_V4.0", torch_dtype=torch.float16, variant="fp16", controlnet=controlnet
).to(device)
pipe.scheduler = diffusers.DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
controlnet_out = pipe(
prompt="high quality photo of a sports bike, city",
negative_prompt="",
guidance_scale=6.5,
num_inference_steps=25,
image=depth_image,
controlnet_conditioning_scale=0.7,
control_guidance_end=0.7,
generator=generator,
).images
controlnet_out[0].save("motorcycle_controlnet_out.png")
希望您会发现 Marigold 有助于解决您的下游任务,无论是更广泛的生成工作流程的一部分,还是感知任务,如 3D 重建。
< > GitHub 上的更新