From b696f63655829732c9fa1ee129b8a155cbe1c9ee Mon Sep 17 00:00:00 2001
From: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com>
Date: Tue, 10 Feb 2026 00:04:56 +0530
Subject: [PATCH] Add bounds check for AuraFlow positional embedding indices

When the input latent spatial dimensions exceed the positional
embedding grid size, the centered-crop index calculation produces
negative or out-of-range indices. This causes a fatal CUDA assertion
error that destroys the CUDA context for the entire process.

Add an explicit check that raises a clear ValueError instead of
silently producing invalid indices.
---
 .../models/transformers/auraflow_transformer_2d.py          | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/diffusers/models/transformers/auraflow_transformer_2d.py b/src/diffusers/models/transformers/auraflow_transformer_2d.py
index e3732662e408..98fa5e3051e9 100644
--- a/src/diffusers/models/transformers/auraflow_transformer_2d.py
+++ b/src/diffusers/models/transformers/auraflow_transformer_2d.py
@@ -76,6 +76,12 @@ def pe_selection_index_based_on_dim(self, h, w):
         h_p, w_p = h // self.patch_size, w // self.patch_size
         h_max, w_max = int(self.pos_embed_max_size**0.5), int(self.pos_embed_max_size**0.5)
 
+        if h_p > h_max or w_p > w_max:
+            raise ValueError(
+                f"Input latent size ({h_p}x{w_p} patches) exceeds the positional embedding grid "
+                f"({h_max}x{w_max}). Use a smaller resolution or increase pos_embed_max_size."
+            )
+
         # Calculate the top-left corner indices for the centered patch grid
         starth = h_max // 2 - h_p // 2
         startw = w_max // 2 - w_p // 2