From baa48d0754d69562ae0d642e2e6ec3a1043e3419 Mon Sep 17 00:00:00 2001
From: rockerBOO <rockerboo@gmail.com>
Date: Thu, 8 May 2025 19:12:10 -0400
Subject: [PATCH] Add Diffusion4k stride to partitioned_vae

---
 flux_train_network.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/flux_train_network.py b/flux_train_network.py
index 1b6c5533..4d99c5d0 100644
--- a/flux_train_network.py
+++ b/flux_train_network.py
@@ -144,6 +144,7 @@ class FluxNetworkTrainer(train_network.NetworkTrainer):
         # Apply partitioned for Diffusion4k
         if args.partitioned_vae:
             ae.decoder.partitioned = True
+            ae.decoder.stride = 2 # Diffusion4k stride
 
         return flux_utils.MODEL_VERSION_FLUX_V1, [clip_l, t5xxl], ae, model