From 1bb7734d105e1925363f6ac10222919e836575b4 Mon Sep 17 00:00:00 2001 From: hlky Date: Wed, 4 Feb 2026 10:59:10 +0000 Subject: [PATCH] ZImageControlNet cfg --- .../pipelines/z_image/pipeline_z_image_controlnet.py | 4 +++- .../pipelines/z_image/pipeline_z_image_controlnet_inpaint.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) 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, )