From a2daa870074310ba2415da993016f0779c8b56e2 Mon Sep 17 00:00:00 2001
From: rockerBOO <rockerboo@gmail.com>
Date: Fri, 28 Feb 2025 14:22:39 -0500
Subject: [PATCH] Add block swap for uncond (neg) for sample images

---
 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 0be81df9..933a4eda 100644
--- a/library/lumina_train_util.py
+++ b/library/lumina_train_util.py
@@ -665,6 +665,7 @@ def denoise(
 
         # compute whether to apply classifier-free guidance based on current timestep
         if current_timestep[0] < cfg_trunc_ratio:
+            model.prepare_block_swap_before_forward()
             noise_pred_uncond = model(
                 img,
                 current_timestep,