up

Author: sayakpaul

src/diffusers/__init__.py

@@ -389,8 +389,10 @@
             "FluxModularPipeline",
             "QwenImageAutoBlocks",
             "QwenImageEditAutoBlocks",
+            "QwenImageEditPlusAutoBlocks",
             "QwenImageEditModularPipeline",
             "QwenImageModularPipeline",
+            "QwenImageEditPlusModularPipeline",
             "StableDiffusionXLAutoBlocks",
             "StableDiffusionXLModularPipeline",
             "WanAutoBlocks",
@@ -1051,8 +1053,10 @@
             FluxModularPipeline,
             QwenImageAutoBlocks,
             QwenImageEditAutoBlocks,
+            QwenImageEditPlusAutoBlocks,
             QwenImageEditModularPipeline,
             QwenImageModularPipeline,
+            QwenImageEditPlusModularPipeline,
             StableDiffusionXLAutoBlocks,
             StableDiffusionXLModularPipeline,
             WanAutoBlocks,

src/diffusers/models/transformers/transformer_qwenimage.py

@@ -134,6 +134,7 @@ def apply_rotary_emb_qwen(
 
         return out
     else:
+        print(f"{x.shape=}, {freqs_cis.shape=}")
         x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
         freqs_cis = freqs_cis.unsqueeze(1)
         x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)

src/diffusers/modular_pipelines/__init__.py

@@ -52,6 +52,8 @@
         "QwenImageModularPipeline",
         "QwenImageEditModularPipeline",
         "QwenImageEditAutoBlocks",
+        "QwenImageEditPlusModularPipeline",
+        "QwenImageEditPlusAutoBlocks",
     ]
     _import_structure["components_manager"] = ["ComponentsManager"]
 
@@ -79,6 +81,8 @@
             QwenImageEditAutoBlocks,
             QwenImageEditModularPipeline,
             QwenImageModularPipeline,
+            QwenImageEditPlusModularPipeline,
+            QwenImageEditPlusAutoBlocks,
         )
         from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline
         from .wan import WanAutoBlocks, WanModularPipeline

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -59,6 +59,7 @@
         ("flux", "FluxModularPipeline"),
         ("qwenimage", "QwenImageModularPipeline"),
         ("qwenimage-edit", "QwenImageEditModularPipeline"),
+        ("qwenimage-edit-plus", "QwenImageEditPlusModularPipeline")
     ]
 )
 
@@ -1628,7 +1629,8 @@ def from_pretrained(
             blocks = ModularPipelineBlocks.from_pretrained(
                 pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
             )
-        except EnvironmentError:
+        except EnvironmentError as e:
+            logger.debug(f"EnvironmentError: {e}")
             blocks = None
 
         cache_dir = kwargs.pop("cache_dir", None)

src/diffusers/modular_pipelines/qwenimage/__init__.py

@@ -38,7 +38,7 @@
         "QwenImageEditAutoBlocks",
         "QwenImageEditPlusAutoBlocks",
     ]
-    _import_structure["modular_pipeline"] = ["QwenImageEditModularPipeline", "QwenImageModularPipeline"]
+    _import_structure["modular_pipeline"] = ["QwenImageEditModularPipeline", "QwenImageModularPipeline", "QwenImageEditPlusModularPipeline"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:

src/diffusers/modular_pipelines/qwenimage/denoise.py

@@ -343,6 +343,7 @@ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState
             cond_kwargs = {k: v for k, v in cond_kwargs.items() if k in guider_input_fields}
 
             # YiYi TODO: add cache context
+            print(f"{block_state.img_shapes=}")
             guider_state_batch.noise_pred = components.transformer(
                 hidden_states=block_state.latent_model_input,
                 timestep=block_state.timestep / 1000,

src/diffusers/modular_pipelines/qwenimage/encoders.py

@@ -132,7 +132,7 @@ def get_qwen_prompt_embeds_edit_plus(
     text_encoder,
     processor,
     prompt: Union[str, List[str]] = None,
-    image: Optional[Union[torch.Tensor, List[PIL.Image.Image], [PIL.Image.Image]]] = None,
+    image: Optional[Union[torch.Tensor, List[PIL.Image.Image], PIL.Image.Image]] = None,
     prompt_template_encode: str = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n",
     img_template_encode: str = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>",
     prompt_template_encode_start_idx: int = 64,
@@ -371,7 +371,7 @@ def intermediate_outputs(self) -> List[OutputParam]:
             ),
             OutputParam(
                 name=self._resized_image_vae_output_name,
-                type_hint=List[PIL.Image.Image],
+                type_hint=torch.Tensor,
                 description="The resized images to be used by the VAE encoder.",
             ),
             OutputParam(
@@ -409,8 +409,8 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
             )
             vae_width, vae_height, _ = calculate_dimensions(self.vae_image_size, image_width / image_height)
             vae_image_sizes.append((vae_width, vae_height))
-            condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
-            vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
+            condition_images.append(components.image_resize_processor.resize(img, condition_height, condition_width))
+            vae_images.append(components.image_resize_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
 
         setattr(block_state, self._resized_image_output_name, condition_images)
         setattr(block_state, self._resized_image_vae_output_name, vae_images)
@@ -718,7 +718,7 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
 
         if components.requires_unconditional_embeds:
             negative_prompt = block_state.negative_prompt or " "
-            block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = get_qwen_prompt_embeds_edit(
+            block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = get_qwen_prompt_embeds_edit_plus(
                 components.text_encoder,
                 components.processor,
                 prompt=negative_prompt,

src/diffusers/pipelines/auto_pipeline.py

@@ -98,6 +98,7 @@
     QwenImageImg2ImgPipeline,
     QwenImageInpaintPipeline,
     QwenImagePipeline,
+    QwenImageEditPlusPipeline,
 )
 from .sana import SanaPipeline
 from .stable_cascade import StableCascadeCombinedPipeline, StableCascadeDecoderPipeline
@@ -186,6 +187,7 @@
         ("flux-kontext", FluxKontextPipeline),
         ("qwenimage", QwenImageImg2ImgPipeline),
         ("qwenimage-edit", QwenImageEditPipeline),
+        ("qwenimage-edit-plus", QwenImageEditPlusPipeline),
     ]
 )