From a77ba5a089b317ebf3b1241ae0c0b85c684431e4 Mon Sep 17 00:00:00 2001 From: Daniel Verdu Date: Fri, 18 Jul 2025 12:49:18 +0200 Subject: [PATCH] fix: Guidance incorrect shape --- .../diffusion_models/flux_kontext/flux_kontext.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py index 49a79b9b..e9eeee57 100644 --- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py +++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py @@ -291,7 +291,8 @@ class FluxKontextModel(BaseModel): else: guidance = torch.tensor( [guidance_embedding_scale], device=self.device_torch) - guidance = guidance.expand(latent_model_input.shape[0]) + # Expand guidance to match original batch_size + guidance = guidance.expand(bs) else: guidance = None