From 9af36fabe70c01221eb39f550393e0421e96594e Mon Sep 17 00:00:00 2001
From: sdbds <865105819@qq.com>
Date: Thu, 5 Sep 2024 14:26:07 +0800
Subject: [PATCH] Update flux_train_utils.py

---
 library/flux_train_utils.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py
index 17126878..1e80f68a 100644
--- a/library/flux_train_utils.py
+++ b/library/flux_train_utils.py
@@ -412,9 +412,10 @@ def get_noisy_model_input_and_timesteps(
         logits_norm = torch.randn(bsz, device=device)
         logits_norm = logits_norm * args.sigmoid_scale  # larger scale for more uniform sampling
         timesteps = logits_norm.sigmoid()
+        x1, x2 = args.min_bucket_reso or min(args.resolution), args.max_bucket_reso or max(args.resolution)
         mu = get_lin_function(
-            y1=get_lin_function(args.min_bucket_reso or min(args.resolution)),
-            y2=get_lin_function(args.max_bucket_reso or max(args.resolution)),
+            y1=x1 / 256 * get_lin_function(x1),
+            y2=4096 / x2 * get_lin_function(x2),
         )((h // 2) * (w // 2))
         timesteps = time_shift(mu, 1.0, timesteps)