Merge pull request #343 from davertor/fix_kontext_bs

fix: Guidance incorrect shape
2025-07-20 12:00:55 -06:00 · 2025-07-20 12:00:55 -06:00 · b1bff66d52
parent 8610c6ed7f a77ba5a089
commit b1bff66d52
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