Update qwenimage/pipeline_qwenimage_edit_plus.py

qwenimage/pipeline_qwenimage_edit_plus.py

@@ -44,7 +44,7 @@ EXAMPLE_DOC_STRING = """
     Examples:
         ```py
         >>> import torch
-        >>> from PIL import Image, ImageOps
+        >>> from PIL import Image
         >>> from diffusers import QwenImageEditPlusPipeline
         >>> from diffusers.utils import load_image
 
@@ -67,23 +67,25 @@ CONDITION_IMAGE_SIZE = 384 * 384
 VAE_IMAGE_SIZE = 1024 * 1024
 
 
+
 def pad_to_aspect(img: Image.Image, target_w: int, target_h: int) -> Image.Image:
-    """Pad (letterbox) to target aspect ratio without warping, then return padded image."""
+    """Pad (letterbox) to target aspect ratio without warping."""
     return ImageOps.pad(
-        img,
+        img.convert("RGB"),
         (int(target_w), int(target_h)),
         method=Image.Resampling.LANCZOS,
         color=(0, 0, 0),
         centering=(0.5, 0.5),
     )
 
+
 def choose_condition_area(canvas_area: int, base_area: int = CONDITION_IMAGE_SIZE) -> int:
     """Choose a conditioning target area derived from canvas area with sensible bounds."""
-    # Scale conditioning area roughly with canvas, anchored at 384^2 when canvas is ~1MP.
     scaled = int(canvas_area * (base_area / (1024 * 1024)))
     return int(min(base_area, max(256 * 256, scaled)))
 
 
+
 # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift
 def calculate_shift(
     image_seq_len,
@@ -685,42 +687,40 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
 
         device = self._execution_device
         # 3. Preprocess image
+
         if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
             if not isinstance(image, list):
                 image = [image]
+
+            # Conditioning resolution derived from canvas area (or overridden)
+            canvas_area = int(width) * int(height)
+            cond_area = int(condition_area) if condition_area is not None else choose_condition_area(canvas_area)
+            cond_w, cond_h = calculate_dimensions(cond_area, float(width) / float(height))
+
             condition_image_sizes = []
             condition_images = []
             vae_image_sizes = []
             vae_images = []
 
-            # Decide which images participate in the VAE latent stream.
-            # Default: all images. This can be overridden (e.g. extras conditioning-only).
+            # Which images participate in the VAE latent stream (default: all)
             if vae_image_indices is None:
                 vae_image_indices = list(range(len(image)))
             vae_set = set(int(i) for i in vae_image_indices)
 
-            # Conditioning resolution: derived from canvas unless explicitly provided.
-            canvas_area = int(width) * int(height)
-            cond_area = int(condition_area) if condition_area is not None else choose_condition_area(canvas_area)
-            cond_w, cond_h = calculate_dimensions(cond_area, float(width) / float(height))
-
             for idx, img in enumerate(image):
-                # Always ensure RGB PIL
-                if isinstance(img, Image.Image):
-                    pil = img.convert("RGB")
-                else:
-                    pil = img
+                # Ensure PIL RGB for padding/resize stability
+                pil = img.convert("RGB") if isinstance(img, Image.Image) else img
 
+                # Strong recommendation: pad to canvas aspect to avoid warping
                 if pad_to_canvas and isinstance(pil, Image.Image):
                     pil = pad_to_aspect(pil, int(width), int(height))
 
-                # Conditioning (VL) path: always include
+                # Conditioning (VL) path: always include, using a canvas-derived size
                 condition_image_sizes.append((cond_w, cond_h))
                 condition_images.append(self.image_processor.resize(pil, cond_h, cond_w))
 
-                # VAE latent path: include only selected indices
+                # VAE path: include only selected indices, and use the *canvas* size
                 if idx in vae_set:
-                    # IMPORTANT: VAE preprocessing uses the *canvas* size to avoid drift/zoom.
                     vae_image_sizes.append((int(width), int(height)))
                     vae_images.append(self.image_processor.preprocess(pil, int(height), int(width)).unsqueeze(2))