fix: Guidance incorrect shape

2025-07-18 12:49:18 +02:00 · 2025-07-18 12:49:18 +02:00 · a77ba5a089
parent 8610c6ed7f
commit a77ba5a089
1 changed files with 2 additions and 1 deletions
--- 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