Update flux_train_utils.py

This commit is contained in:
sdbds
2024-09-05 14:26:07 +08:00
parent 406e739cb9
commit 9af36fabe7

View File

@@ -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)