Diffusers documentation

Shap-E

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Shap-E

Shap-E 模型由 OpenAI 的 Alex Nichol 和 Heewoo Jun 在 Shap-E: Generating Conditional 3D Implicit Functions 中提出。

该论文的摘要如下:

我们提出了 Shap-E,一种用于 3D 资产的条件生成模型。与近期生成单一输出表示的 3D 生成模型不同,Shap-E 直接生成隐式函数的参数,这些参数既可以渲染为带纹理的网格,也可以渲染为神经辐射场。我们在两个阶段训练 Shap-E:首先,训练一个编码器,将 3D 资产确定性地映射为隐式函数的参数;其次,在编码器的输出上训练一个条件扩散模型。当在大规模的 3D 与文本配对数据集上进行训练时,我们得到的模型能够在几秒钟内生成复杂且多样的 3D 资产。与基于点云的显式生成模型 Point-E 相比,Shap-E 收敛更快,并且在建模更高维、多表示的输出空间的情况下,达到了相当甚至更好的样本质量。

原始代码库可在 openai/shap-e 中找到。

请参阅跨流水线复用组件部分,了解如何高效地将相同组件加载到多个流水线中。

请确保已安装以下库。

# uncomment to install the necessary libraries in Colab
#!pip install -q diffusers transformers accelerate trimesh

文本到 3D

要生成 3D 对象的 gif,请将文本提示传递给 ShapEPipeline。该流水线会生成一系列图像帧,用于创建 3D 对象。

import torch
from diffusers import ShapEPipeline

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

pipe = ShapEPipeline.from_pretrained("openai/shap-e", dtype=torch.float16, variant="fp16")
pipe = pipe.to(device)

guidance_scale = 15.0
prompt = ["A firecracker", "A birthday cupcake"]

images = pipe(
    prompt,
    guidance_scale=guidance_scale,
    num_inference_steps=64,
    frame_size=256,
).images

现在使用 export_to_gif() 函数将图像帧列表转换为 3D 对象的 gif 文件。

from diffusers.utils import export_to_gif

export_to_gif(images[0], "firecracker_3d.gif")
export_to_gif(images[1], "cake_3d.gif")
prompt = "A firecracker"
prompt = "一个生日纸杯蛋糕"

图像到 3D

要从另一张图像生成 3D 对象,请使用 ShapEImg2ImgPipeline。您可以使用现有图像,也可以生成一张全新的图像。让我们使用 Kandinsky 2.1 模型来生成一张新图像。

from diffusers import DiffusionPipeline
import torch

prior_pipeline = DiffusionPipeline.from_pretrained("kandinsky-community/kandinsky-2-1-prior", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
pipeline = DiffusionPipeline.from_pretrained("kandinsky-community/kandinsky-2-1", dtype=torch.float16, use_safetensors=True).to("cuda")

prompt = "A cheeseburger, white background"

image_embeds, negative_image_embeds = prior_pipeline(prompt, guidance_scale=1.0).to_tuple()
image = pipeline(
    prompt,
    image_embeds=image_embeds,
    negative_image_embeds=negative_image_embeds,
).images[0]

image.save("burger.png")

将芝士汉堡传递给 ShapEImg2ImgPipeline,以生成其 3D 表示。

from PIL import Image
from diffusers import ShapEImg2ImgPipeline
from diffusers.utils import export_to_gif

pipe = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", dtype=torch.float16, variant="fp16").to("cuda")  # or "mps", "xpu", "cpu"

guidance_scale = 3.0
image = Image.open("burger.png").resize((256, 256))

images = pipe(
    image,
    guidance_scale=guidance_scale,
    num_inference_steps=64,
    frame_size=256,
).images

gif_path = export_to_gif(images[0], "burger_3d.gif")
cheeseburger
3D 芝士汉堡

生成网格

Shap-E 是一个灵活的模型,还可以生成带纹理的网格输出,以便渲染用于下游应用。在本示例中,我们将输出转换为 glb 文件,因为 🤗 Datasets 库支持对 glb 文件进行网格可视化,这些文件可由 Dataset viewer 渲染。

通过将 output_type 参数指定为 "mesh",您可以为 ShapEPipeline 和 ShapEImg2ImgPipeline 生成网格输出:

import torch
from diffusers import ShapEPipeline

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

pipe = ShapEPipeline.from_pretrained("openai/shap-e", dtype=torch.float16, variant="fp16")
pipe = pipe.to(device)

guidance_scale = 15.0
prompt = "A birthday cupcake"

images = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=64, frame_size=256, output_type="mesh").images

使用 export_to_ply() 函数将网格输出保存为 ply 文件:

您可以选择使用 export_to_obj() 函数将网格输出保存为 obj 文件。支持以多种格式保存网格输出,使其在下游应用中更加灵活!

from diffusers.utils import export_to_ply

ply_path = export_to_ply(images[0], "3d_cake.ply")
print(f"Saved to folder: {ply_path}")

然后,你可以使用 trimesh 库将 ply 文件转换为 glb 文件:

import trimesh

mesh = trimesh.load("3d_cake.ply")
mesh_export = mesh.export("3d_cake.glb", file_type="glb")

默认情况下,网格输出聚焦于底部视角,但你可以通过应用旋转变换来更改默认视角:

import trimesh
import numpy as np

mesh = trimesh.load("3d_cake.ply")
rot = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])
mesh = mesh.apply_transform(rot)
mesh_export = mesh.export("3d_cake.glb", file_type="glb")

将网格文件上传到您的数据集仓库,以便使用数据集查看器进行可视化!

ShapEPipeline

class diffusers.ShapEPipeline

< >

( prior: PriorTransformertext_encoder: CLIPTextModelWithProjectiontokenizer: CLIPTokenizerscheduler: HeunDiscreteSchedulershap_e_renderer: ShapERenderer )

参数

  • prior (PriorTransformer) — 用于从文本嵌入近似图像嵌入的 unCLIP 标准先验模型。
  • text_encoder (CLIPTextModelWithProjection) — 冻结的文本编码器。
  • tokenizer (CLIPTokenizer) — 用于对文本进行分词的 CLIPTokenizer。
  • scheduler (HeunDiscreteScheduler) — 用于与 prior 模型配合以生成图像嵌入的调度器。
  • shap_e_renderer (ShapERenderer) — Shap-E 渲染器将生成的潜变量投影为 MLP 的参数,并使用 NeRF 渲染方法创建 3D 对象。

用于生成 3D 资产潜在表示并使用 NeRF 方法进行渲染的 Pipeline。

该模型继承自 DiffusionPipeline。有关为所有 pipeline 实现的通用方法(下载、保存、在特定设备上运行等),请参阅父类文档。

__call__

< >

( prompt: strnum_images_per_prompt: int = 1num_inference_steps: int = 25generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneguidance_scale: float = 4.0frame_size: int = 64output_type: str | None = 'pil'return_dict: bool = True ) → ShapEPipelineOutput or tuple

参数

  • prompt (str 或 list[str]) — 用于引导图像生成的提示词或提示词列表。
  • num_images_per_prompt (int,可选,默认为 1) — 每个提示词生成的图像数量。
  • num_inference_steps (int,可选,默认为 25) — 去噪步数。更多的去噪步数通常能生成更高质量的图像,但代价是推理速度变慢。
  • generator (torch.Generator 或 list[torch.Generator],可选) — 用于使生成过程确定性的 torch.Generator。
  • latents (torch.Tensor, 可选) — 预先生成的、从高斯分布中采样的含噪潜变量,用作图像生成的输入。可用于以不同的提示词微调同一生成过程。如果未提供,则使用提供的随机generator进行采样,生成一个潜变量张量。
  • guidance_scale (float,可选,默认值为 4.0) — 较高的 guidance scale 值会促使模型生成与文本 prompt 紧密相关的图像,但代价是图像质量降低。当 guidance_scale > 1 时,guidance scale 生效。
  • frame_size (int, 可选,默认为 64) — 生成的 3D 输出中每个图像帧的宽度和高度。
  • output_type (str,可选,默认为 "pil") — 生成图像的输出格式。可选 "pil"(PIL.Image.Image)、"np" (np.array)、"latent"(torch.Tensor)或 mesh(MeshDecoderOutput)。
  • return_dict (bool,可选,默认为 True) — 是否返回 ShapEPipelineOutput 而不是普通的元组。

返回值

ShapEPipelineOutput 或 tuple

如果 return_dict 为 True,则返回 ShapEPipelineOutput,否则返回一个 tuple,其中第一个元素是包含生成图像的列表。

用于生成的 pipeline 的调用函数。

示例:

>>> import torch
>>> from diffusers import DiffusionPipeline
>>> from diffusers.utils import export_to_gif

>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

>>> repo = "openai/shap-e"
>>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)
>>> pipe = pipe.to(device)

>>> guidance_scale = 15.0
>>> prompt = "a shark"

>>> images = pipe(
...     prompt,
...     guidance_scale=guidance_scale,
...     num_inference_steps=64,
...     frame_size=256,
... ).images

>>> gif_path = export_to_gif(images[0], "shark_3d.gif")

ShapEImg2ImgPipeline

class diffusers.ShapEImg2ImgPipeline

< >

( prior: PriorTransformerimage_encoder: CLIPVisionModelimage_processor: CLIPImageProcessorPilscheduler: HeunDiscreteSchedulershap_e_renderer: ShapERenderer )

参数

  • prior (PriorTransformer) — 用于从文本嵌入近似图像嵌入的 unCLIP 标准先验模型。
  • image_encoder (CLIPVisionModel) — 冻结的图像编码器。
  • image_processor (CLIPImageProcessor) — 用于处理图像的 CLIPImageProcessor。
  • scheduler (HeunDiscreteScheduler) — 用于与 prior 模型配合以生成图像嵌入的调度器。
  • shap_e_renderer (ShapERenderer) — Shap-E 渲染器将生成的潜变量投影为 MLP 的参数,并使用 NeRF 渲染方法创建 3D 对象。

用于从图像生成 3D 资产的潜在表示,并使用 NeRF 方法进行渲染的 Pipeline。

该模型继承自 DiffusionPipeline。有关为所有 pipeline 实现的通用方法(下载、保存、在特定设备上运行等),请参阅父类文档。

__call__

< >

( image: PIL.Image.Image | list[PIL.Image.Image]num_images_per_prompt: int = 1num_inference_steps: int = 25generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneguidance_scale: float = 4.0frame_size: int = 64output_type: str | None = 'pil'return_dict: bool = True ) → ShapEPipelineOutput or tuple

参数

  • image (torch.Tensor、PIL.Image.Image、np.ndarray、list[torch.Tensor]、list[PIL.Image.Image] 或 list[np.ndarray]) — Image 或表示图像批次的张量,用作起始点。也可以接受图像潜变量作为 image,但如果直接传入潜变量,则不会再次进行编码。
  • num_images_per_prompt (int,可选,默认为 1) — 每个提示词生成的图像数量。
  • num_inference_steps (int,可选,默认为 25) — 去噪步数。更多的去噪步数通常能生成更高质量的图像,但代价是推理速度变慢。
  • generator (torch.Generator 或 list[torch.Generator],可选) — 用于使生成过程确定性的 torch.Generator。
  • latents (torch.Tensor, 可选) — 预先生成的、从高斯分布中采样的含噪潜变量,用作图像生成的输入。可用于以不同的提示词微调同一生成过程。如果未提供,则使用提供的随机generator进行采样,生成一个潜变量张量。
  • guidance_scale (float,可选,默认值为 4.0) — 较高的 guidance scale 值会促使模型生成与文本 prompt 紧密相关的图像,但代价是图像质量降低。当 guidance_scale > 1 时,guidance scale 生效。
  • frame_size (int, 可选,默认为 64) — 生成的 3D 输出中每个图像帧的宽度和高度。
  • output_type (str,可选,默认为 "pil") — 生成图像的输出格式。可选 "pil"(PIL.Image.Image)、"np"(np.array)、"latent"(torch.Tensor)或 mesh(MeshDecoderOutput)。
  • return_dict (bool,可选,默认为 True) — 是否返回 ShapEPipelineOutput 而不是普通元组。

返回值

ShapEPipelineOutput 或 tuple

如果 return_dict 为 True,则返回 ShapEPipelineOutput,否则返回一个 tuple,其中第一个元素是包含生成图像的列表。

用于生成的 pipeline 的调用函数。

示例:

>>> from PIL import Image
>>> import torch
>>> from diffusers import DiffusionPipeline
>>> from diffusers.utils import export_to_gif, load_image

>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

>>> repo = "openai/shap-e-img2img"
>>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)
>>> pipe = pipe.to(device)

>>> guidance_scale = 3.0
>>> image_url = "https://hf.co/datasets/diffusers/docs-images/resolve/main/shap-e/corgi.png"
>>> image = load_image(image_url).convert("RGB")

>>> images = pipe(
...     image,
...     guidance_scale=guidance_scale,
...     num_inference_steps=64,
...     frame_size=256,
... ).images

>>> gif_path = export_to_gif(images[0], "corgi_3d.gif")

ShapEPipelineOutput

class diffusers.pipelines.shap_e.pipeline_shap_e.ShapEPipelineOutput

< >

( images: list[list[PIL.Image.Image]] | list[list[numpy.ndarray]] )

参数

  • images (torch.Tensor) — 用于 3D 渲染的图像列表。

ShapEPipeline 和 ShapEImg2ImgPipeline 的输出类。

Update on GitHub