diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py index 08fc4da0e7ba..3c8db4a0f748 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py @@ -635,10 +635,12 @@ def __call__( latent_model_input = latents_typed.repeat(2, 1, 1, 1) prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds timestep_model_input = timestep.repeat(2) + control_image_input = control_image.repeat(2, 1, 1, 1, 1) else: latent_model_input = latents.to(self.transformer.dtype) prompt_embeds_model_input = prompt_embeds timestep_model_input = timestep + control_image_input = control_image latent_model_input = latent_model_input.unsqueeze(2) latent_model_input_list = list(latent_model_input.unbind(dim=0)) @@ -647,7 +649,7 @@ def __call__( latent_model_input_list, timestep_model_input, prompt_embeds_model_input, - control_image, + control_image_input, conditioning_scale=controlnet_conditioning_scale, ) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py index 3b0f8dc288d3..cdc60eaf4dd3 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py @@ -657,10 +657,12 @@ def __call__( latent_model_input = latents_typed.repeat(2, 1, 1, 1) prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds timestep_model_input = timestep.repeat(2) + control_image_input = control_image.repeat(2, 1, 1, 1, 1) else: latent_model_input = latents.to(self.transformer.dtype) prompt_embeds_model_input = prompt_embeds timestep_model_input = timestep + control_image_input = control_image latent_model_input = latent_model_input.unsqueeze(2) latent_model_input_list = list(latent_model_input.unbind(dim=0)) @@ -669,7 +671,7 @@ def __call__( latent_model_input_list, timestep_model_input, prompt_embeds_model_input, - control_image, + control_image_input, conditioning_scale=controlnet_conditioning_scale, )