Package: cf-guidance
cf_guidanceThis library has normalizations and schedules for the Classifier-free
Guidance parameter in Diffusion models. We call the guidance modifierstransforms.It seems thattransformscan improve the quality of generated images
across Diffusion pipelines. To be specific:Normalizing can help with image syntax.Scheduling can improve image details.For more details, please check out the officialblog
post.Installpipinstallcf_guidanceHow to useThe follow section creates a few example schedules. Then, we use the
schedules to initialize different kinds oftransforms.Creating Cosine schedulesWe create a few different schedules:A Cosine schedule that starts from the default Guidance parameter in
diffusion $\left( G = 7.5\right)$ and works it way down to a minimum of
$0.15$.Next we create the same schedule, except with an initialwarmup. With
warmup, $G$ starts from $0.15$ and linearly works up to a maximum
$G_\text{max} = 7.5$. Here we warmup over the first 10% of the total
diffusion steps. After warmup, the schedule decreases following a cosine
curve.Finally, we create a schedule with small values of $G$. This is meant
for the guidance transformTNormGuidance.
This transforms puts the vectors on a similar scale, so we need a much
smaller value than usual when mixing them.fromcf_guidance.schedulesimportget_cos_sched# Parameters from the blog postmax_val=7.5min_val=0.15num_steps=50num_warmup_steps=5# 1) Baseline cosine schedulecos_params={'max_val':max_val,'num_steps':num_steps,'min_val':min_val,'num_warmup_steps':0,}# 2) Cosine schedule with warmupwarmup_cos_params={'max_val':max_val,'num_steps':num_steps,'min_val':min_val,'warmup_init_val':min_val+1.,# to show we can offset the warmup relative to min'num_warmup_steps':num_warmup_steps,}# 3) Cosine schedule with small values for T-Norm guidancet_scale_params={'max_val':0.25,'num_steps':num_steps,'min_val':0.05,}# create the schedulescos_g=get_cos_sched(**cos_params)warmup_g=get_cos_sched(**warmup_cos_params)TNorm_g=get_cos_sched(**t_scale_params)Let’s plot these schedules to see what they look like.importmatplotlib.pyplotasplt# plot themplt.plot(cos_g)plt.xlabel('Diffusion Timesteps')plt.ylabel('$G$ Guidance Parameter')plt.title('Cosine Schedule');plt.plot(warmup_g)plt.xlabel('Diffusion Timesteps')plt.ylabel('$G$ Guidance Parameter')plt.title('Warmup Cosine Schedule');plt.plot(TNorm_g)plt.xlabel('Diffusion Timesteps')plt.ylabel('$G$ Guidance Parameter')plt.title('Small Cosine Schedule to T-Norm');Creating Guidance TransformsGuidance transforms take one initialization parameter:schedules. This
is a map from parameter names to an array-like, indexable sequence of
values.For a given parameternameat diffusion timestepidx, the value ofschedules[name][idx]should be the value of that parameter at the
given timestep.In this case, we call the parameter $G$ as a lowercase ‘g’.fromcf_guidance.transformsimport*# Create a baseline Guidance with cosine schedule.sample_sched={'g':cos_g}guidance=GuidanceTfm(sample_sched)# Guidance with normalization and a cosine schedule.norm_guidance=BaseNormGuidance(sample_sched)# T-Norm Guidance with the smaller-valued schedulesmall_sched={'g':TNorm_g}small_guidance=TNormGuidance(small_sched)Using the transforms in a Diffusion pipelineThe code snippet below shows where and how these transforms are used in
an imagegenerate()pipeline. We will use thenorm_guidanceclass
created above for the example.Specifically, we call thisnorm_tfmwith the following arguments:The unconditioned noise predictions.The conditional noise predictions.The current timestep index.We borrow code from theofficialStableDiffusionPipelinefrom HuggingFace to show where ournorm_tfmgoes. This seems like a
good starting point, since many scripts and functions are based on this
HuggingFace script.# inside of `StableDiffusionPipeline`fori,tinenumerate(self.progress_bar(timesteps_tensor)):# expand the latents if we are doing classifier free guidancelatent_model_input=torch.cat([latents]*2)ifdo_classifier_free_guidanceelselatentslatent_model_input=self.scheduler.scale_model_input(latent_model_input,t)# predict the noise residualnoise_pred=self.unet(latent_model_input,t,encoder_hidden_states=text_embeddings).sample# NOTE: our transforms go here:###############################ifdo_classifier_free_guidance:noise_pred_uncond,noise_pred_text=noise_pred.chunk(2)## OLD UPADTE#noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)# NEW cf_guidance UPDATEnoise_pred=norm_guidance(noise_pred_uncond,noise_pred_text,i)Creating more cosine schedulesThe cosine schedule function was based on a combination of the Cosine
schedulers intimmandHuggingFace.It has a variety of parameters to support many different schedule
combinations, as shown below.# using wilder parameters, worth exploringtwo_cycles={'max_val':max_val,'num_steps':num_steps,'min_val':min_val,'num_cycles':1,'warmup_init_val':min_val+1.,# to show we can offset the warmup relative to min'num_warmup_steps':0,}k_decay_cos={'max_val':max_val,'num_steps':num_steps,'min_val':min_val,'num_cycles':1.5,'k_decay':0.7,'warmup_init_val':min_val+1.,# to show we can offset the warmup relative to min'num_warmup_steps':num_warmup_steps,}# create the schedulestwo_cycle_sched=get_cos_sched(**two_cycles)k_decay_sched=get_cos_sched(**k_decay_cos)plt.plot(two_cycle_sched)plt.title('Cosine With a Full Cycle');plt.plot(k_decay_sched)plt.title('Cosine with Warmup, 1.5 Cycles, and K-decay');