[SDXL] Make sure multi batch prompt embeds works

[patrickvonplaten]

What does this PR do?

Currently the following does not work:

import hf_image_uploader as hiu
import torch
from compel import Compel, ReturnedEmbeddingsType
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
  "stabilityai/stable-diffusion-xl-base-1.0",
  variant="fp16",
  torch_dtype=torch.float16
)
pipe2 = DiffusionPipeline.from_pretrained(
  "stabilityai/stable-diffusion-xl-refiner-1.0",
  variant="fp16",
  torch_dtype=torch.float16
)
pipe.enable_model_cpu_offload()
pipe2.enable_model_cpu_offload()

compel = Compel(
  tokenizer=[pipe.tokenizer, pipe.tokenizer_2] ,
  text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
  returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
  requires_pooled=[False, True]
)

compel2 = Compel(
  tokenizer=pipe.tokenizer_2,
  text_encoder=pipe.text_encoder_2,
  returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
  requires_pooled=True,
)

# apply weights
prompt = ["a red cat playing with a (ball)1.5", "a red cat playing with a (ball)0.6"]
conditioning, pooled = compel(prompt)
conditioning2, pooled2 = compel2(prompt)

# generate image

for _ in range(3):
    generator = [torch.Generator().manual_seed(i) for i in range(len(prompt))]

    image = pipe(prompt_embeds=conditioning, pooled_prompt_embeds=pooled, generator=generator, num_inference_steps=30, output_type="latent").images
    image = pipe2(image=image, prompt_embeds=conditioning2, pooled_prompt_embeds=pooled2, generator=generator, num_inference_steps=20).images[0]
    hiu.upload(image, "patrickvonplaten/images")

This PR fixes it.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[DN6]

LGTM 👍🏽 . Just had one small comment/question.

[bghira]

this fixes the issue i described a while back, thank you