diff --git a/src/flux_1_schnell/latent_creator/latent_creator.py b/src/flux_1_schnell/latent_creator/latent_creator.py
index 82309ee..aa251db 100644
--- a/src/flux_1_schnell/latent_creator/latent_creator.py
+++ b/src/flux_1_schnell/latent_creator/latent_creator.py
@@ -5,7 +5,7 @@ class LatentCreator:
 
     @staticmethod
     def create(height: int, width: int, seed: int) -> (mx.array, mx.array):
-        r_h = height // 16
-        r_w = width // 16
-        latents = mx.random.normal(shape=[1, r_h * r_w, 64], key=mx.random.key(seed))
+        latent_height = height // 16
+        latent_width = width // 16
+        latents = mx.random.normal(shape=[1, latent_height * latent_width, 64], key=mx.random.key(seed))
         return latents
diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py
index e0d9434..7676885 100644
--- a/src/flux_1_schnell/models/transformer/transformer.py
+++ b/src/flux_1_schnell/models/transformer/transformer.py
@@ -67,14 +67,14 @@ class Transformer(nn.Module):
         return noise
 
     @staticmethod
-    def _prepare_latent_image_ids(width, height) -> mx.array:
-        r_h = height // 16
-        r_w = width // 16
-        latent_image_ids = mx.zeros((r_h, r_w, 3))
-        latent_image_ids = latent_image_ids.at[:, :, 1].add(mx.arange(0, r_h)[:, None])
-        latent_image_ids = latent_image_ids.at[:, :, 2].add(mx.arange(0, r_w)[None, :])
+    def _prepare_latent_image_ids(width: int, height: int) -> mx.array:
+        latent_height = height // 16
+        latent_width = width // 16
+        latent_image_ids = mx.zeros((latent_height, latent_width, 3))
+        latent_image_ids = latent_image_ids.at[:, :, 1].add(mx.arange(0, latent_height)[:, None])
+        latent_image_ids = latent_image_ids.at[:, :, 2].add(mx.arange(0, latent_width)[None, :])
         latent_image_ids = mx.repeat(latent_image_ids[None, :], 1, axis=0)
-        latent_image_ids = mx.reshape(latent_image_ids, (1, r_h * r_w, 3))
+        latent_image_ids = mx.reshape(latent_image_ids, (1, latent_height * latent_width, 3))
         return latent_image_ids
 
     @staticmethod