diffusers icon indicating copy to clipboard operation
diffusers copied to clipboard

`StateManager`, `Hook`, `State`: missing context setting (cond or uncond) in pipelines

Open toilaluan opened this issue 2 months ago • 2 comments

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(...)

toilaluan avatar Dec 01 '25 02:12 toilaluan

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

toilaluan avatar Dec 01 '25 03:12 toilaluan

@DN6 @toilaluan can you review my PR

omkar-334 avatar Dec 14 '25 13:12 omkar-334