faster_cfg.py

#Put it on stable-diffusion-webui/scripts
#Idea:https://twitter.com/Birchlabs/status/1640033271512702977
#Related to https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/9129
import torch

import modules.scripts as scripts
import gradio as gr

from modules.script_callbacks import CFGDenoiserParams, on_cfg_denoiser,  CFGDenoisedParams, on_cfg_denoised
from modules.processing import StableDiffusionProcessing, process_images
import modules.shared as shared

class Script(scripts.Script):

	def __init__(self):
		pass

	def title(self):
		return "faster_cfg"

	def ui(self, is_img2img):
		with gr.Group():
			with gr.Row():
				enable = gr.Checkbox(value=False, label="Enable")
				sigma = gr.Slider(minimum=0, maximum=3, step=0.05, label="Threshold of sigma", value=1.1)
					
		return enable, sigma
	
	def denoiser_callback(self, params: CFGDenoiserParams):
		if self.enable:
			if params.sigma[0].item() < self.sigma:
				# only x for cond
				params.x = params.x[:params.x.shape[0] // 2]
				shared.batch_cond_uncond = False
				if not self.skip_cfg:
					self.skip_cfg = True
					print("skipping cfg !!! ")

	def denoised_callback(self, params: CFGDenoisedParams):
		if self.enable:
			if self.skip_cfg:
				# copy x_cond
				# x_cond == x_cond + scale * (x_cond - x_cond)
				params.x = torch.cat([params.x] * 2)
				shared.batch_cond_uncond = self.batch_cond_uncond	
									 
	def run(
		self, 
		p: StableDiffusionProcessing, 
		enable:bool, 
		sigma: float, 
		):
		
		self.enable = enable
		self.sigma = sigma
		self.batch_cond_uncond = shared.batch_cond_uncond
		self.skip_cfg = False
		
		if not hasattr(self, 'callbacks_added'):
			on_cfg_denoiser(self.denoiser_callback)
			on_cfg_denoised(self.denoised_callback)
			self.callbacks_added = True

		return process_images(p)