`StateManager`, `Hook`, `State`: missing context setting (cond or uncond) in pipelines
In diffusion denoising model, we usually use classifier free guidance to control qualities. With State Manager, we want each context hold different state, so we have to set which context when calling self.transformer, here is a good example
But some popular pipelines (Flux2) or pipeline variants (Kontext, Image to image, Inpainting) don't implement this. For example, Flux T2I has it while Flux Kontext is missing
A Fix should be simple but require us to check all pipelines
with self.transformer.cache_context("cond"):
noise_pred = self.transformer(...)
with self.transformer.cache_context("uncond"):
neg_noise_pred = self.transformer(...)
By using a simple script that counts self.transformer( and self.transformer.cache_context occurrences, there are pipelines need to be fixed:
src/diffusers/pipelines/omnigen/pipeline_omnigen.py
src/diffusers/pipelines/kandinsky/text_encoder.py
src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py
src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py
src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py
src/diffusers/pipelines/unidiffuser/modeling_uvit.py
src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py
src/diffusers/pipelines/chroma/pipeline_chroma.py
src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py
src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py
src/diffusers/pipelines/flux/pipeline_flux_controlnet.py
src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py
src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py
src/diffusers/pipelines/flux/pipeline_flux_fill.py
src/diffusers/pipelines/flux/pipeline_flux_kontext.py
src/diffusers/pipelines/flux/pipeline_flux_inpaint.py
src/diffusers/pipelines/flux/pipeline_flux_img2img.py
src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py
src/diffusers/pipelines/flux/pipeline_flux_control.py
src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py
src/diffusers/pipelines/bria/pipeline_bria.py
src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py
src/diffusers/pipelines/flux2/pipeline_flux2.py
src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py
src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py
src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py
src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py
src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py
src/diffusers/pipelines/sana/pipeline_sana_sprint.py
src/diffusers/pipelines/sana/pipeline_sana.py
src/diffusers/pipelines/sana/pipeline_sana_controlnet.py
src/diffusers/pipelines/chronoedit/pipeline_chronoedit.py
src/diffusers/pipelines/consisid/pipeline_consisid.py
src/diffusers/pipelines/amused/pipeline_amused_inpaint.py
src/diffusers/pipelines/amused/pipeline_amused_img2img.py
src/diffusers/pipelines/amused/pipeline_amused.py
src/diffusers/pipelines/lumina/pipeline_lumina.py
src/diffusers/pipelines/cogview3/pipeline_cogview3plus.py
src/diffusers/pipelines/cogview4/pipeline_cogview4_control.py
src/diffusers/pipelines/easyanimate/pipeline_easyanimate.py
src/diffusers/pipelines/easyanimate/pipeline_easyanimate_inpaint.py
src/diffusers/pipelines/easyanimate/pipeline_easyanimate_control.py
src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py
src/diffusers/pipelines/latte/pipeline_latte.py
src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py
src/diffusers/pipelines/sana_video/pipeline_sana_video.py
src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py
src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py
src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py
src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_i2v.py
src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2.py
src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_v2v.py
src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing_i2v.py
src/diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py
src/diffusers/pipelines/kandinsky5/pipeline_kandinsky.py
src/diffusers/pipelines/wan/pipeline_wan_video2video.py
src/diffusers/pipelines/z_image/pipeline_z_image.py
src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py
src/diffusers/pipelines/dit/pipeline_dit.py
src/diffusers/pipelines/prx/pipeline_prx.py
src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py
src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py
src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py
src/diffusers/pipelines/lumina2/pipeline_lumina2.py
src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py
src/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py
src/diffusers/pipelines/allegro/pipeline_allegro.py
src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py
src/diffusers/pipelines/pag/pipeline_pag_sd_3.py
src/diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py
src/diffusers/pipelines/pag/pipeline_pag_sana.py
src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py
src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py
Check Script
#!/bin/bash
# Usage: ./script.sh [directory]
# If no directory is provided, it uses the current directory.
directory="${1:-.}"
find "$directory" -type f | while read -r file; do
N=$(grep -I -o -F "self.transformer(" "$file" | wc -l)
K=$(grep -I -o -F "self.transformer.cache_context(" "$file" | wc -l)
if [ "$N" -ne "$K" ]; then
echo "$file"
fi
done
@DN6 @toilaluan can you review my PR