Stable Diffusion and SDXL Callbacks are fundamentally broken for prompt_embeds

[AI-Casanova]

Describe the bug

I'm working on implementing scheduled prompting in SDNext and realized that the callback in SDXL pipelines is fundamentally not functional. This also applies to Stable Diffusion Pipelines and likely others.

diffusers/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

Lines 1245 to 1252 in dac623b

	prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
	negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
	add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
	negative_pooled_prompt_embeds = callback_outputs.pop(
	"negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
	)
	add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
	negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids)

Of the preceding lines the following variables are completely useless as they are not referenced inside the denoising loop and thus adjusting them in the callback can do nothing:

• negative_prompt_embeds
• negative_pooled_prompt_embeds
• negative_add_time_ids

If self.do_classifier_free_guidance is True, the other variable are concatenations of themselves and their negative counterparts, and not what their variables state.

diffusers/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

Line 1126 in dac623b

# 7. Prepare added time ids & embeddings

For the callback to operate intuitively, the code following # 7 should be split out into a separate function, and following the variables NOT overwritten:

• prompt_embeds
• add_text_embeds (naming convention should be fixed)
• add_time_ids

That function should then be called post-callback to update all of the associated concat variables.

Reproduction

N/A

This is a code only issue.

Logs

No response

System Info

N/A

Who can help?

@yiyixuxu @sayakpaul

[AI-Casanova]

I can and will create a monkeypatch in my code for the time being, but that is a lot of code duplication for something that should be handled inside the pipelines.

(OK the monkeypatch isn't that bad, but the problem remains)

[yiyixuxu]

can you provide the callback example that you're trying to implement?
it is true that the negative* embeddings are not used inside the denoising loop but it will be passed to your callback function and you will be able to use it inside

[vladmandic]

the point of the callback is to be two way so values can be changed. if some values are two way and some are one way only, it's bad.

what we're trying to do is to alter emdeds per step.

[AI-Casanova]

The problem comes when attempting to feed prompt embeddings back into the denoising loop, here's an attempt at non-optimized code:

def callback_fn(pipe, i, t, callback_kwargs):
    if i > 10:  # arbitrary set point
        prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt("Prompt for last half of denoising")
        callback_kwargs["prompt_embeds"] = prompt_embeds
        callback_kwargs["negative_prompt_embeds"] = negative_prompt_embeds
        callback_kwargs["add_text_embeds"] = pooled_prompt_embeds
        callback_kwargs["negative_pooled_prompt_embeds"] = negative_pooled_prompt_embeds
    return callback_kwargs

This works if CFG is disabled, because the negatives are ignored

If CFG is enabled:

        callback_kwargs["prompt_embeds"] = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
        callback_kwargs["add_text_embeds"] = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)

works, but again fails if CFG is disabled, so the callback is required to determine

[AI-Casanova]

for reference

Switching between a woman sitting on a rotten log and a woman sitting on a park bench
20 steps total switching at the indicated step