undefined variable in PAG pipeline_pag_sd

[christopher5106]

Describe the bug

noise_pred_text is not defined here in the case of do_perturbed_attention_guidance=True, guidance_rescale >0 and do_classifier_free_guidance=True

Reproduction

Done

Logs

No response

System Info

Linux

Who can help?

No response

[yiyixuxu]

thanks!
would you be willing to open a PR to help us fix it?

[christopher5106]

I can, but I'm not a big fan of these if statements that are hard to read, you see, prone to errors.

Instead of

                if self.do_classifier_free_guidance:
                    negative_image_embeds, image_embeds = image_embeds.chunk(2)

                if self.do_perturbed_attention_guidance:
                    image_embeds = self._prepare_perturbed_attention_guidance(
                        image_embeds, negative_image_embeds, self.do_classifier_free_guidance
                    )
                elif self.do_classifier_free_guidance:
                    image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0)

wouldn't i t be more readable to have

                negative_image_embeds = None
                if self.do_classifier_free_guidance:
                    negative_image_embeds, image_embeds = image_embeds.chunk(2)
                    
                if self.do_perturbed_attention_guidance:
                    image_embeds = torch.cat([image_embeds] * 2, dim=0)

                if self.do_classifier_free_guidance:
                    image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0)

and not use the method _prepare_perturbed_attention_guidance()

Or even shorter:

                if self.do_perturbed_attention_guidance:
                    image_embeds = torch.cat([
                            image_embeds, 
                            image_embeds.chunk(2)[1] if self.do_classifier_free_guidance else image_embeds
                    ] , dim=0)

The last snippet has the advantage to not interfere with the main flow when PAG not active.

[christopher5106]

The same for:

# perform guidance
                if self.do_perturbed_attention_guidance:
                        if self.do_classifier_free_guidance:
                            noise_pred_uncond, noise_pred_text, noise_pred_perturb = noise_pred_batch.chunk(3)

                    noise_pred = self._apply_perturbed_attention_guidance(
                        noise_pred, self.do_classifier_free_guidance, self.guidance_scale, t
                    )

                elif self.do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)

                if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)

To me it's hard to read. I would prefer:

# perform guidance
if self.do_perturbed_attention_guidance and self.do_classifier_free_guidance:
     noise_pred_uncond, noise_pred_text, noise_pred_perturb = noise_pred.chunk(3)
elif self.do_perturbed_attention_guidance:
    noise_pred_text, noise_pred_perturb = noise_pred.chunk(2)
elif self.do_classifier_free_guidance:
   noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)

if self.do_perturbed_attention_guidance or self.do_classifier_free_guidance:
   noise_pred = noise_pred_text
if self.do_perturbed_attention_guidance:
   pag_scale = self._get_pag_scale(t) 
   noise_pred += pag_scale * (noise_pred_text - noise_pred_perturb)
if self.do_classifier_free_guidance:
   noise_pred += (guidance_scale - 1) * (noise_pred_text - noise_pred_uncond)
if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
   # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
   noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)

[yiyixuxu]

Hi @christopher5106
100% agree with you that it is less readable; we are also not a fan of adding additional methods like _apply_perturbed_attention_guidance as it is not necessary and hurts the readability

However, this was a conscious trade-off made for PAG pipelines; PAG works great with almost all of our pipelines, so, potentially, we could have one PAG pipeline for each of the existing pipelines we already have. We are a small team so maintenance burden becomes an obvious concern, and we want to be able work with the community as much as possible; therefore PAG pipelines are written in this way with a goal that it should be very easy for us to integrate and maintain and to work with our community. we pretty much want to guarantee:

1. all the PAG-related logic is inside if self.do_perturbed_attention_guidance statement and nowhere else in the pipeline.__call__ and extracted away in PAGMixin
2. all the pipeline methods are copied from the base pipeline with #Copied from statement so that it is easier for us to make sure PAG is up to date with its base pipeline

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[yiyixuxu]

closing as it is fixed in #9537