From acfc448f0c51744661d329ed6822f7d9c1be5856 Mon Sep 17 00:00:00 2001 From: Fabio Date: Sun, 18 Aug 2024 17:56:31 +0200 Subject: [PATCH] fix guidance scale --- src/flux_1_schnell/models/transformer/transformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py index 5c1b2b7..fa4e356 100644 --- a/src/flux_1_schnell/models/transformer/transformer.py +++ b/src/flux_1_schnell/models/transformer/transformer.py @@ -37,7 +37,7 @@ class Transformer(nn.Module): time_step = config.sigmas[t] * config.num_train_steps time_step = mx.broadcast_to(time_step, (1,)) hidden_states = self.x_embedder(hidden_states) - guidance = mx.broadcast_to(config.guidance, (1,)) + guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)) text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance) encoder_hidden_states = self.context_embedder(prompt_embeds) txt_ids = Transformer._prepare_text_ids(seq_len = prompt_embeds.shape[1])