From f0731883004dd1738d03c3fcfbd71c1d6bc0c9c9 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 22 Sep 2026 12:14:19 -0600 Subject: [PATCH] Fix Qwen Image Edit Plus batch_size > 1 without cached text embeddings by splitting the batch-stacked control tensor into one control image per prompt. --- .../qwen_image/qwen_image_edit_plus.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py index 8f4068f..b4ea52e 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py @@ -176,8 +176,15 @@ class QwenImageEditPlusModel(QwenImageModel): if control_images is None: raise ValueError("Missing control images for QwenImageEditPlusModel") - - if not isinstance(control_images, list): + + if isinstance(control_images, torch.Tensor): + # training with one control image per item hands us the batch-stacked + # tensor (bs, ch, h, w); one set per prompt, not one set of bs images + if control_images.dim() == 4 and control_images.shape[0] == len(prompt): + control_images = [[img] for img in control_images] + else: + control_images = [control_images] + elif not isinstance(control_images, list): control_images = [control_images] # expects a list of list of control images List[List[Tensor]] where each item corresponds to a batch item,