在最终图像生成StableDiffusionPipeline.pretrained之前获取图像迭代

Question

我目前正在使用diffusers StableDiffusionPipeline (从拥抱脸)，以产生一个不和谐的机器人，我使用与我的朋友的AI图像。我想知道是否有可能在图像完成之前获得生成图像的预览？

例如，如果一幅图像需要20秒才能生成，因为它使用的是扩散，所以它从模糊开始，并逐渐变得越来越好。我想要的是将图像保存在每次迭代(或每隔几秒钟)上，并查看其进展情况。我怎么能这么做？

class ImageGenerator:
    def __init__(self, socket_listener, pretty_logger, prisma):
        self.model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", revision="fp16", torch_dtype=torch.float16, use_auth_token=os.environ.get("HF_AUTH_TOKEN"))
        self.model = self.model.to("cuda")

    async def generate_image(self, data):
        start_time = time.time()
        with autocast("cuda"):
            image = self.model(data.description, height=self.default_height, width=self.default_width,
                               num_inference_steps=self.default_inference_steps, guidance_scale=self.default_guidance_scale)
            image.save(...)

我目前的代码是这样的，但是它只在图像完全完成时才返回。我试图了解StableDiffusionPipeline内部是如何生成图像的，但是我在任何地方都找不到生成图像的地方。如果有人能提供任何关于我可以从哪里开始的指示/提示，那将是非常有帮助的。

Answer 1

您可以使用稳定扩散管道的回调参数来获得图像：链接到文档的潜在空间表示。

实现显示了如何将潜在项转换回图像。我们只需复制这些代码并解码那些潜在的内容。

下面是一个小示例，它每5个步骤保存生成的映像：

from diffusers import StableDiffusionPipeline
import torch

#load model
model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", revision="fp16", torch_dtype=torch.float16, use_auth_token="YOUR TOKEN HERE")
model = model.to("cuda")

def callback(iter, t, latents):
    # convert latents to image
    with torch.no_grad():
        latents = 1 / 0.18215 * latents
        image = model.vae.decode(latents).sample

        image = (image / 2 + 0.5).clamp(0, 1)

        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
        image = image.cpu().permute(0, 2, 3, 1).float().numpy()

        # convert to PIL Images
        image = model.numpy_to_pil(image)

        # do something with the Images
        for i, img in enumerate(image):
            img.save(f"iter_{iter}_img{i}.png")

# generate image (note the `callback` and `callback_steps` argument)
image = model("tree", callback=callback, callback_steps=5)

为了理解稳定的扩散模型，我强烈推荐这博客文章。