在生成最终图像之前获取图像迭代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)

要了解稳定扩散模型，我强烈推荐this博客文章。