fix qwenimage mask ordering, tensor images, layered output, vae clamp, and exports

src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py

@@ -1028,6 +1028,7 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutp
             result_rows.append(torch.cat(result_row, dim=-1))
 
         dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
+        dec = torch.clamp(dec, min=-1.0, max=1.0)
 
         if not return_dict:
             return (dec,)

src/diffusers/models/transformers/transformer_qwenimage.py

@@ -897,6 +897,8 @@ def forward(
 
         if self.zero_cond_t:
             timestep = torch.cat([timestep, timestep * 0], dim=0)
+            if additional_t_cond is not None:
+                additional_t_cond = torch.cat([additional_t_cond, additional_t_cond], dim=0)
             modulate_index = torch.tensor(
                 [[0] * prod(sample[0]) + [1] * sum([prod(s) for s in sample[1:]]) for sample in img_shapes],
                 device=timestep.device,

src/diffusers/modular_pipelines/qwenimage/encoders.py

@@ -520,7 +520,10 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
         resized_images = []
         resized_cond_images = []
         for image in images:
-            image_width, image_height = image.size
+            if isinstance(image, torch.Tensor):
+                image_width, image_height = image.shape[-1], image.shape[-2]
+            else:
+                image_width, image_height = image.size
 
             # For VAE encoder (1024x1024 target area)
             vae_width, vae_height, _ = calculate_dimensions(1024 * 1024, image_width / image_height)

src/diffusers/modular_pipelines/qwenimage/inputs.py

@@ -220,32 +220,19 @@ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -
         block_state.batch_size = block_state.prompt_embeds.shape[0]
         block_state.dtype = block_state.prompt_embeds.dtype
 
-        _, seq_len, _ = block_state.prompt_embeds.shape
-
-        block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
-        block_state.prompt_embeds = block_state.prompt_embeds.view(
-            block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+        block_state.prompt_embeds = block_state.prompt_embeds.repeat_interleave(
+            block_state.num_images_per_prompt, dim=0
         )
-
-        block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.repeat(1, block_state.num_images_per_prompt, 1)
-        block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.view(
-            block_state.batch_size * block_state.num_images_per_prompt, seq_len
+        block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.repeat_interleave(
+            block_state.num_images_per_prompt, dim=0
         )
 
         if block_state.negative_prompt_embeds is not None:
-            _, seq_len, _ = block_state.negative_prompt_embeds.shape
-            block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
-                1, block_state.num_images_per_prompt, 1
-            )
-            block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
-                block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
-            )
-
-            block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.repeat(
-                1, block_state.num_images_per_prompt, 1
+            block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat_interleave(
+                block_state.num_images_per_prompt, dim=0
             )
-            block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.view(
-                block_state.batch_size * block_state.num_images_per_prompt, seq_len
+            block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.repeat_interleave(
+                block_state.num_images_per_prompt, dim=0
             )
 
         self.set_block_state(state, block_state)

src/diffusers/pipelines/qwenimage/__init__.py

@@ -12,7 +12,7 @@
 
 _dummy_objects = {}
 _additional_imports = {}
-_import_structure = {"pipeline_output": ["QwenImagePipelineOutput", "QwenImagePriorReduxPipelineOutput"]}
+_import_structure = {"pipeline_output": ["QwenImagePipelineOutput"]}
 
 try:
     if not (is_transformers_available() and is_torch_available()):
@@ -22,7 +22,6 @@
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
-    _import_structure["modeling_qwenimage"] = ["ReduxImageEncoder"]
     _import_structure["pipeline_qwenimage"] = ["QwenImagePipeline"]
     _import_structure["pipeline_qwenimage_controlnet"] = ["QwenImageControlNetPipeline"]
     _import_structure["pipeline_qwenimage_controlnet_inpaint"] = ["QwenImageControlNetInpaintPipeline"]

src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py

@@ -247,20 +247,15 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device)
 
         prompt_embeds = prompt_embeds[:, :max_sequence_length]
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
             prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py

@@ -312,20 +312,15 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device)
 
         prompt_embeds = prompt_embeds[:, :max_sequence_length]
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
             prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py

@@ -294,20 +294,15 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device)
 
         prompt_embeds = prompt_embeds[:, :max_sequence_length]
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
             prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py

@@ -297,18 +297,13 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device)
 
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None
@@ -620,7 +615,11 @@ def __call__(
             [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
             returning a tuple, the first element is a list with the generated images.
         """
-        image_size = image[0].size if isinstance(image, list) else image.size
+        img_for_size = image[0] if isinstance(image, list) else image
+        if isinstance(img_for_size, torch.Tensor):
+            image_size = (img_for_size.shape[-1], img_for_size.shape[-2])
+        else:
+            image_size = img_for_size.size
         calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
         height = height or calculated_height
         width = width or calculated_width

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py

@@ -309,18 +309,13 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device)
 
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None
@@ -780,7 +775,11 @@ def __call__(
             [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
             returning a tuple, the first element is a list with the generated images.
         """
-        image_size = image[0].size if isinstance(image, list) else image.size
+        img_for_size = image[0] if isinstance(image, list) else image
+        if isinstance(img_for_size, torch.Tensor):
+            image_size = (img_for_size.shape[-1], img_for_size.shape[-2])
+        else:
+            image_size = img_for_size.size
         calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
 
         # height and width are the same as the calculated height and width

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py

@@ -311,18 +311,13 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device)
 
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None
@@ -642,7 +637,11 @@ def __call__(
             [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
             returning a tuple, the first element is a list with the generated images.
         """
-        image_size = image[-1].size if isinstance(image, list) else image.size
+        img_for_size = image[-1] if isinstance(image, list) else image
+        if isinstance(img_for_size, torch.Tensor):
+            image_size = (img_for_size.shape[-1], img_for_size.shape[-2])
+        else:
+            image_size = img_for_size.size
         calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
         height = height or calculated_height
         width = width or calculated_width

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py

@@ -290,20 +290,15 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device)
 
         prompt_embeds = prompt_embeds[:, :max_sequence_length]
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
             prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py

@@ -301,20 +301,15 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device)
 
         prompt_embeds = prompt_embeds[:, :max_sequence_length]
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
             prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py

@@ -313,20 +313,15 @@ def encode_prompt(
         device = device or self._execution_device
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
         if prompt_embeds is None:
             prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, device)
 
         prompt_embeds = prompt_embeds[:, :max_sequence_length]
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
 
         if prompt_embeds_mask is not None:
             prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
-            prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+            prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
 
             if prompt_embeds_mask.all():
                 prompt_embeds_mask = None
@@ -659,7 +654,11 @@ def __call__(
             [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
             returning a tuple, the first element is a list with the generated images.
         """
-        image_size = image[0].size if isinstance(image, list) else image.size
+        img_for_size = image[0] if isinstance(image, list) else image
+        if isinstance(img_for_size, torch.Tensor):
+            image_size = (img_for_size.shape[-1], img_for_size.shape[-2])
+        else:
+            image_size = img_for_size.size
         assert resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {resolution}"
         calculated_width, calculated_height = calculate_dimensions(
             resolution * resolution, image_size[0] / image_size[1]
@@ -875,7 +874,7 @@ def __call__(
 
         self._current_timestep = None
         if output_type == "latent":
-            image = latents
+            images = latents
         else:
             latents = self._unpack_latents(latents, height, width, layers, self.vae_scale_factor)
             latents = latents.to(self.vae.dtype)