几行代码实现可视化的文生图采样过程

前言

Huggingface的StableDiffusionPipeline实现了端到端的文生图功能，用户输入文字，即可得到最后的图片。但是在某些场景下，如果需要获取中间过程的图片，用于优化长时间等待过程的用户体验，就只能通过修改源码进行实现了。

源码分析

对于SD的__call__()方法来说，前面的所有步骤都是准备工作，关键在于第七步denoising这里。progress_bar里面是一个循环，timesteps用于控制采样步数，因此总共循环timesteps步，而progress_bar之外是图片尺寸提升以及安全特征检测等图像恢复工作。不难发现，latents 是采样过程中的潜在空空间向量，需要通过VAE进行解码才能变为现实空间的图片向量。所以为了将采样过程中间生成的图片也显示出来，只需要把最后出现的VAE也用于中间生成的潜在空间向量的解码即可。

# diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

	# 7. Denoising loop
	num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
	self._num_timesteps = len(timesteps)
	with self.progress_bar(total=num_inference_steps) as progress_bar:
	    for i, t in enumerate(timesteps):
	        # expand the latents if we are doing classifier free guidance
	        latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
	        latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
	
	        # predict the noise residual
	        noise_pred = self.unet(
	            latent_model_input,
	            t,
	            encoder_hidden_states=prompt_embeds,
	            timestep_cond=timestep_cond,
	            cross_attention_kwargs=self.cross_attention_kwargs,
	            return_dict=False,
	        )[0]
	
	        # perform guidance
	        if self.do_classifier_free_guidance:
	            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
	            noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
	
	        if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
	            # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
	            noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)
	
	        # compute the previous noisy sample x_t -> x_t-1
	        latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
	        self.mid_res.append(latents)    
	        if callback_on_step_end is not None:
	            callback_kwargs = {}
	            for k in callback_on_step_end_tensor_inputs:
	                callback_kwargs[k] = locals()[k]
	            callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
	
	            latents = callback_outputs.pop("latents", latents)
	            prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
	            negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
	
	        # call the callback, if provided
	        if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
	            progress_bar.update()
	            if callback is not None and i % callback_steps == 0:
	                step_idx = i // getattr(self.scheduler, "order", 1)
	                callback(step_idx, t, latents)
	
	if not output_type == "latent":
	    # 图像从潜在空间转化为图片空间，4通道转3通道，并且图片尺寸变大：从64提升到512
	    image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[0]
	    # 图片特征安全性检测
	    image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
	else:
	    image = latents
	    has_nsfw_concept = None
	
	if has_nsfw_concept is None:
	    do_denormalize = [True] * image.shape[0]
	else:
	    do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
	# tensor转化为PIL.image
	image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
	
	# Offload all models
	self.maybe_free_model_hooks()
	
	if not return_dict:
	    return (image, has_nsfw_concept)
	
	return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

修改

如果每一采样的中间过程都进行解码，那势必增加整体推理时间，因此可以跳过采样初期，直到采样后期图片基本确定后，再按照一定的间隔进行解码。
i<20表示采样步骤小于20不考虑， i%2==0表示后期每两步进行一次解码，加入yield变成生成器。

if i < 20 or i%2==0:
    continue
mid_img = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[0]
yield self.image_processor.postprocess(mid_img, output_type=output_type, do_denormalize=[True] * mid_img.shape[0])

在原代码相应位置修改后结果如下：

# diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
	with self.progress_bar(total=num_inference_steps) as progress_bar:
		for i, t in enumerate(timesteps):
			 # ...省略中间代码...
             # call the callback, if provided
             if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                 progress_bar.update()
                 if callback is not None and i % callback_steps == 0:
                     step_idx = i // getattr(self.scheduler, "order", 1)
                     callback(step_idx, t, latents)
             # 加入这一部分代码
             if i < 20 or i%2==0:
                 continue
             mid_img = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[0]
             yield self.image_processor.postprocess(mid_img, output_type=output_type, do_denormalize=[True] * mid_img.shape[0])
	
	if not output_type == "latent":
	 		# ...省略后面代码...

效果

通过gradio构建一个可视化应用：

# 导入相关的模块
from diffusers import StableDiffusionPipeline
import torch
import time
import gradio as gr
# 本地模型路径
model_id = "your_model_path"
pipe = StableDiffusionPipeline.from_pretrained(model_id)

# 定义一个事件监听函数，返回生成器对象
def stream_img(prompt):
    for mid in pipe(prompt=prompt, num_inference_steps=30):
        yield mid[0]

with gr.Blocks(analytics_enabled=False) as demo:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(lines=2, placeholder="Prompt (press Ctrl+Enter or Alt+Enter to generate)")
            button = gr.Button('生成', size='sm')
        output_img = gr.Image(interactive=False, type="pil",  image_mode="RGB")
        # 输出到图片组件
        button.click(fn=stream_img, inputs=prompt, outputs=output_img, queue=True)
demo.queue(max_size=20)
demo.launch()

需要注意：button.click(queue=True)才能接受生成器对象，并且需要设置demo.queue(max_size)

---

可以看到图片生成的中间过程正确显示出来了。