CombinedTimestepGuidanceTextProjEmbeddings.forward() missing 1 required positional argument: 'pooled_projection'

[christopher5106]

Describe the bug

There is a bug when using controlnet_union with flux schnell.

The reason is because controlnet union uses guidance_embeds set to True while Schnell doesn't.

Reproduction

The code in the pipeline should be written that way I believe:

                # handle guidance
                if self.controlnet.config.guidance_embeds:
                    guidance = torch.tensor([guidance_scale], device=device)
                    guidance = guidance.expand(latents.shape[0])
                else:
                    guidance = None

                # controlnet
                controlnet_block_samples, controlnet_single_block_samples = self.controlnet(
                    hidden_states=latents,
                    controlnet_cond=control_image,
                    controlnet_mode=control_mode,
                    conditioning_scale=controlnet_conditioning_scale,
                    timestep=timestep / 1000,
                    guidance=guidance,
                    pooled_projections=pooled_prompt_embeds,
                    encoder_hidden_states=prompt_embeds,
                    txt_ids=text_ids,
                    img_ids=latent_image_ids,
                    joint_attention_kwargs=self.joint_attention_kwargs,
                    return_dict=False,
                )

                if self.transformer.config.guidance_embeds:
                    guidance = torch.tensor([guidance_scale], device=device)
                    guidance = guidance.expand(latents.shape[0])
                else:
                    guidance = None

                noise_pred = self.transformer(
                    hidden_states=latents,
                    timestep=timestep / 1000,
                    guidance=guidance,
                    pooled_projections=pooled_prompt_embeds,
                    encoder_hidden_states=prompt_embeds,
                    controlnet_block_samples=controlnet_block_samples,
                    controlnet_single_block_samples=controlnet_single_block_samples,
                    txt_ids=text_ids,
                    img_ids=latent_image_ids,
                    joint_attention_kwargs=self.joint_attention_kwargs,
                    return_dict=False,
                )[0]

instead of this

One is conditioned on self.controlnet.config.guidance_embeds while flux is conditioned on self.transformer.config.guidance_embeds. I didn't test it.

Logs

No response

System Info

Ubuntu

Who can help?

No response

[christopher5106]

@sayakpaul @yiyixuxu @DN6

[sayakpaul]

But I don't think reasonable to add the schnell support if the underlying ControlNet wasn't trained on Schnell.

[christopher5106]

Arrrh if we start to only follow ML guidelines about training distributions, then probably 99% of our users are wrong about mixing things that have never been guaranteed to work... mixing a controlnet with an ip adapter is 100% safe ? that means did we train them together ? we now in a world of "out of training distribution" all the time, if it works.

It's not a big deal to read the conf of the controlnet to infer the controls, and to read the conf of the model to make the full inference.

[christopher5106]

I can try to deliver you an empirical study if it works, that might be a good starting point.

[christopher5106]

I run the flux-schnell with controlnet union thanks to a hack to have the pipeline work (loading the controlnet without guidance embed, which is not its default config, should not be like that), otherwise following this example :
prompt: 'A bohemian-style female travel blogger with sun-kissed skin and messy beach waves.'
guidance: 3
controlnet strength: 0.5
canny control image:

results:

[yiyixuxu]

thanks @christopher5106
looking into this!

[culda]

Just adding onto this, I'm getting this error on flux dev, not just schnell.

[christopher5106]

Check the value of guidance_embeds in your flux dev transformer config, and in the controlnet config. If they are not the same, then it's due to the same error.