Fix: QwenImageEditPlusModel crashes with batch_size > 1
Pass only the current sample's caption to encode_prompt, and right-pad per-sample embeddings before concatenating.
This commit is contained in:
1 parent
b1bf3e4381
commit
b90cccdfb1
1 file changed
+15
-1
@@ -211,7 +211,7 @@ class QwenImageEditPlusModel(QwenImageModel):
|
||||
)
|
||||
|
||||
prompt_embeds, prompt_embeds_mask = self.pipeline.encode_prompt(
|
||||
prompt,
|
||||
[prompt[b]],
|
||||
image=batch_control_images,
|
||||
device=self.device_torch,
|
||||
num_images_per_prompt=1,
|
||||
@@ -223,6 +223,20 @@ class QwenImageEditPlusModel(QwenImageModel):
|
||||
)
|
||||
prompt_embeds_list.append(prompt_embeds)
|
||||
prompt_embeds_mask_list.append(prompt_embeds_mask)
|
||||
|
||||
# each sample is encoded separately, so sequence lengths differ.
|
||||
# pad on the right before concatenating, same as diffusers does internally.
|
||||
max_seq_len = max(e.shape[1] for e in prompt_embeds_list)
|
||||
for i, (e, m) in enumerate(zip(prompt_embeds_list, prompt_embeds_mask_list)):
|
||||
pad_len = max_seq_len - e.shape[1]
|
||||
if pad_len > 0:
|
||||
prompt_embeds_list[i] = torch.cat(
|
||||
[e, e.new_zeros(e.shape[0], pad_len, e.shape[2])], dim=1
|
||||
)
|
||||
prompt_embeds_mask_list[i] = torch.cat(
|
||||
[m, m.new_zeros(m.shape[0], pad_len)], dim=1
|
||||
)
|
||||
|
||||
pe = PromptEmbeds(torch.cat(prompt_embeds_list, dim=0))
|
||||
pe.attention_mask = torch.cat(prompt_embeds_mask_list, dim=0)
|
||||
return pe
|
||||
|
||||
Reference in new issue
Block a user