From 03b8f860a4de2c6b3fcffbffb8e4edd25d368d19 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Wed, 8 Jul 2026 18:07:59 -0700 Subject: [PATCH] fix qwenimage mask ordering, tensor images, layered output, vae clamp, and exports --- .../autoencoders/autoencoder_kl_qwenimage.py | 1 + .../transformers/transformer_qwenimage.py | 2 ++ .../modular_pipelines/qwenimage/encoders.py | 5 +++- .../modular_pipelines/qwenimage/inputs.py | 29 +++++-------------- src/diffusers/pipelines/qwenimage/__init__.py | 3 +- .../pipelines/qwenimage/pipeline_qwenimage.py | 9 ++---- .../pipeline_qwenimage_controlnet.py | 9 ++---- .../pipeline_qwenimage_controlnet_inpaint.py | 9 ++---- .../qwenimage/pipeline_qwenimage_edit.py | 15 +++++----- .../pipeline_qwenimage_edit_inpaint.py | 15 +++++----- .../qwenimage/pipeline_qwenimage_edit_plus.py | 15 +++++----- .../qwenimage/pipeline_qwenimage_img2img.py | 9 ++---- .../qwenimage/pipeline_qwenimage_inpaint.py | 9 ++---- .../qwenimage/pipeline_qwenimage_layered.py | 17 +++++------ 14 files changed, 55 insertions(+), 92 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py index 220520a12e68..28970757f00b 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py +++ b/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,) diff --git a/src/diffusers/models/transformers/transformer_qwenimage.py b/src/diffusers/models/transformers/transformer_qwenimage.py index 464712bd94fd..1e960b3d016c 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage.py +++ b/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, diff --git a/src/diffusers/modular_pipelines/qwenimage/encoders.py b/src/diffusers/modular_pipelines/qwenimage/encoders.py index 5dade5716a49..ca7368163cc6 100644 --- a/src/diffusers/modular_pipelines/qwenimage/encoders.py +++ b/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) diff --git a/src/diffusers/modular_pipelines/qwenimage/inputs.py b/src/diffusers/modular_pipelines/qwenimage/inputs.py index faec7db245df..95dfd8bfa806 100644 --- a/src/diffusers/modular_pipelines/qwenimage/inputs.py +++ b/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) diff --git a/src/diffusers/pipelines/qwenimage/__init__.py b/src/diffusers/pipelines/qwenimage/__init__.py index 3f43d0ebb0b9..1837f4a9e889 100644 --- a/src/diffusers/pipelines/qwenimage/__init__.py +++ b/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"] diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py index 1da0518a4f65..68e795afcac4 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py +++ b/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 diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py index f946fdf27d00..40749e24bbbc 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py +++ b/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 diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py index 97f510a6dbf4..c51fec9dcdb4 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py +++ b/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 diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py index 85abb815cf23..5b83b1c98650 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py +++ b/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 diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py index 57d1fdaaf99f..f49fa53eb12b 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py +++ b/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 diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py index 84d1b60152b1..deeadb3e83bb 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_plus.py +++ b/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 diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py index 9b9af83737e5..9fe6243ca2df 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py +++ b/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 diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py index 3d5f0040932a..859a7170d436 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py +++ b/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 diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py index 7e06a7d36ffd..4f7e52aec5a1 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_layered.py +++ b/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)