From 1bba7acd9ac42ef5a654cadf47356d20d407ce82 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Fri, 28 Feb 2025 14:11:53 -0500 Subject: [PATCH] Add block swap in sample image timestep loop --- library/lumina_train_util.py | 1 + 1 file changed, 1 insertion(+) diff --git a/library/lumina_train_util.py b/library/lumina_train_util.py index e008b3ce..0be81df9 100644 --- a/library/lumina_train_util.py +++ b/library/lumina_train_util.py @@ -647,6 +647,7 @@ def denoise( """ for i, t in enumerate(tqdm(timesteps)): + model.prepare_block_swap_before_forward() # reverse the timestep since Lumina uses t=0 as the noise and t=1 as the image current_timestep = 1 - t / scheduler.config.num_train_timesteps