本文约4800字,建议阅读9分钟本文介绍了Stable Diffusion是一个文本到图像的潜在扩散模型的入门介绍及使用教程。

Stable Diffusion是一个文本到图像的潜在扩散模型,由CompVis、Stability AI和LAION的研究人员和工程师创建。它使用来自LAION-5B数据库子集的512x512图像进行训练。使用这个模型,可以生成包括人脸在内的任何图像,因为有开源的预训练模型,所以我们也可以在自己的机器上运行它,如下图所示。

!nvidia-smi !pip install diffusers==0.4.0 !pip install transformers scipy ftfy !pip install "ipywidgets>=7,<8"
这里需要的额外操作是必须同意模型协议,还要通过勾选复选框来接受模型许可。“Hugging Face”上注册,并获得访问令牌等等。另外对于谷歌collab,它已经禁用了外部小部件,所以需要启用它。运行以下代码这样才能够使用“notebook_login”
from google.colab import output output.enable_custom_widget_manager()
现在就可以从的账户中获得的访问令牌登录Hugging Face了:
from huggingface_hub import notebook_login notebook_login()
从diffusers库加载StableDiffusionPipeline。StableDiffusionPipeline是一个端到端推理管道,可用于从文本生成图像。我们将加载预训练模型权重。模型id将是CompVis/ stable-diffusion-v1-4,我们也将使用一个特定类型的修订版torch_dtype函数。设置revision = “fp16”从半精度分支加载权重,并设置torch_dtype = " torch。torch_dtype = “torch.float16”告诉模型使用fp16的权重。像这样设置可以减少内存,并且运行的更快。
import torch from diffusers import StableDiffusionPipeline # make sure you're logged in with `huggingface-cli login` pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)
下面设置GPU:
pipe = pipe.to("cuda")
现在就可以生成图片了。我们将编写一个提示文本并将其交给管道并打印输出。这里的输入提示是“an astronaut riding a horse”,让看看输出:
prompt = "a photograph of an astronaut riding a horse" image = pipe(prompt).images[0] # image here is in [PIL format](https://pillow.readthedocs.io/en/stable/) # Now to display an image you can do either save it such as: image.save(f"astronaut_rides_horse.png") import torch generator = torch.Generator("cuda").manual_seed(1024) image = pipe(prompt, generator=generator).images[0] image
还可以使用num_inference_steps参数更改步骤的数量。一般来说,推理步骤越多,生成的图像质量越高,但生成结果需要更多的时间。如果你想要更快的结果,你可以使用更少的步骤。下面的单元格使用与前面相同的种子,但步骤更少。注意一些细节,如马头或头盔,比前一张图定义得更模糊:
import torch generator = torch.Generator("cuda").manual_seed(1024) image = pipe(prompt, num_inference_steps=15, generator=generator).images[0] image from PIL import Image def image_grid(imgs, rows, cols): assert len(imgs) == rows*cols w, h = imgs[0].size grid = Image.new('RGB', size=(cols*w, rows*h)) grid_w, grid_h = grid.size for i, img in enumerate(imgs): grid.paste(img, box=(i%cols*w, i//cols*h)) return grid
现在,我们可以生成多个图像并一起展示了。
num_images = 3 prompt = ["a photograph of an astronaut riding a horse"] * num_images images = pipe(prompt).images grid = image_grid(images, rows=1, cols=3) grid num_cols = 3 num_rows = 4 prompt = ["a photograph of an astronaut riding a horse"] * num_cols all_images = [] for i in range(num_rows): images = pipe(prompt).images all_images.extend(images) grid = image_grid(all_images, rows=num_rows, cols=num_cols) grid prompt = "a photograph of an astronaut riding a horse" image = pipe(prompt, height=512, width=768).images[0] image import torch torch_device = "cuda" if torch.cuda.is_available() else "cpu"
预训练的模型包括建立一个完整的管道所需的所有组件。它们存放在以下文件夹中:text_encoder:Stable Diffusion使用CLIP,但其他扩散模型可能使用其他编码器,如BERT。tokenizer:它必须与text_encoder模型使用的标记器匹配。scheduler:用于在训练过程中逐步向图像添加噪声的scheduler算法。U-Net:用于生成输入的潜在表示的模型。VAE,我们将使用它将潜在的表示解码为真实的图像。可以通过引用组件被保存的文件夹,使用from_pretraining的子文件夹参数来加载组件。
from transformers import CLIPTextModel, CLIPTokenizer from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler # 1. Load the autoencoder model which will be used to decode the latents into image space. vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae") # 2. Load the tokenizer and text encoder to tokenize and encode the text. tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") # 3. The UNet model for generating the latents. unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
现在,我们不加载预定义的scheduler,而是加载K-LMS:
from diffusers import LMSDiscreteScheduler scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
将模型移动到GPU上。
vae = vae.to(torch_device) text_encoder = text_encoder.to(torch_device) unet = unet.to(torch_device)
定义用于生成图像的参数。与前面的示例相比,设置num_inference_steps = 100来获得更明确的图像。
prompt = ["a photograph of an astronaut riding a horse"] height = 512 # default height of Stable Diffusion width = 512 # default width of Stable Diffusion num_inference_steps = 100 # Number of denoising steps guidance_scale = 7.5 # Scale for classifier-free guidance generator = torch.manual_seed(32) # Seed generator to create the inital latent noise batch_size = 1
获取文本提示的text_embeddings。然后将嵌入用于调整U-Net模型。
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt") with torch.no_grad(): text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
获得用于无分类器引导的无条件文本嵌入,这只是填充令牌(空文本)的嵌入。它们需要具有与text_embeddings (batch_size和seq_length)相同的形状。
 max_length = text_input.input_ids.shape[-1] uncond_input = tokenizer( [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt" ) with torch.no_grad(): uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
对于无分类的引导,需要进行两次向前传递。第一个是条件输入(text_embeddings),第二个是无条件嵌入(uncond_embeddings)。把两者连接到一个批处理中,以避免进行两次向前传递:
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
生成初始随机噪声:
latents = torch.randn( (batch_size, unet.in_channels, height // 8, width // 8), generator=generator, ) latents = latents.to(torch_device)
产生的形状为64 * 64的随机潜在空间。模型会将这种潜在的表示(纯噪声)转换为512 * 512的图像。使用所选的num_inference_steps初始化scheduler。这将计算sigma和去噪过程中使用的确切步长值:
scheduler.set_timesteps(num_inference_steps)
K-LMS需要用它的sigma值乘以潜在空间的值:

latents = latents * scheduler.init_noise_sigma
最后就是去噪的循环:
from tqdm.auto import tqdm from torch import autocast for t in tqdm(scheduler.timesteps): # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. latent_model_input = torch.cat([latents] * 2) latent_model_input = scheduler.scale_model_input(latent_model_input, t) # predict the noise residual with torch.no_grad(): noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample # perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 latents = scheduler.step(noise_pred, t, latents).prev_sample
然后就是使用vae可将产生的潜在空间解码回图像:
# scale and decode the image latents with vae latents = 1 / 0.18215 * latents with torch.no_grad(): image = vae.decode(latents).sample
最后将图像转换为PIL,以便我们可以显示或保存它。
image = (image / 2 + 0.5).clamp(0, 1) image = image.detach().cpu().permute(0, 2, 3, 1).numpy() images = (image * 255).round().astype("uint8") pil_images = [Image.fromarray(image) for image in images]
pil_images[0]