Diffusers
Published:
Diffusers 是 Diffusion 模型的代码框架包装,其最大的单位是 pipeline。你只需要一个 pipeline 和权重文件(夹),就可以做到几行代码生成图片。例子如下:
from diffusers import DiffusionPipeline
import torch
pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5", torch_dtype=torch.float16, variant='fp16')
pipe.to("cuda")
pipe("A dog").images[0]
现在一些论文(比如 layout-guidance,简称 tflcg)利用 diffusers 库,结合自己的算法提供了一个 pipeline 供大家使用,比如其库提供的示例代码:
from tflcg.layout_guidance_pipeline import LayoutGuidanceStableDiffusionPipeline
from diffusers import EulerDiscreteScheduler
import transformers
import torch
transformers.utils.move_cache()
pipe = LayoutGuidanceStableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5", torch_dtype=torch.float16, variant='fp16')
pipe = pipe.to("mps")
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
prompt = "A cat playing with a ball"
bboxes = [[0.55, 0.4, 0.95, 0.8]]
image = pipe(prompt, num_inference_steps=20,
token_indices=[[2]],
bboxes=bboxes).images[0]
image = pipe.draw_box(image, bboxes)
image.save("output.png")
为了能够在未来的工作中写出类似的代码,我觉得我有必要研究一下 diffusers 库的相关代码细节。我参考的是知乎上大佬的这篇文章并且补充了一些细节。