diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py
index df3950a..5c63d07 100644
--- a/src/flux_1_schnell/config/config.py
+++ b/src/flux_1_schnell/config/config.py
@@ -15,7 +15,7 @@ class Config:
             height: int = 1024,
             guidance: float = 4.0,
     ):
-        if width %16 != 0 or height % 16 != 0:
+        if width % 16 != 0 or height % 16 != 0:
             log.warning("Width and height should be multiples of 16. Rounding down.")
         self.width = 16 * (height // 16)
         self.height = 16 * (width // 16)
diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py
index 053c49b..bbed325 100644
--- a/src/flux_1_schnell/flux.py
+++ b/src/flux_1_schnell/flux.py
@@ -56,12 +56,12 @@ class Flux1:
 
             mx.eval(latents)
 
-        latents = Flux1._unpack_latents(latents, config.width, config.height)
+        latents = Flux1._unpack_latents(latents, config.height, config.width)
         decoded = self.vae.decode(latents)
         return ImageUtil.to_image(decoded)
 
     @staticmethod
-    def _unpack_latents(latents, width, height):
+    def _unpack_latents(latents: mx.array, width: int, height: int) -> mx.array:
         latents = mx.reshape(latents, (1, height//16, width//16, 16, 2, 2))
         latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
         latents = mx.reshape(latents, (1, 16, height//16 *2, width//16 * 2))
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 0074e0d..5c1b2b7 100644
--- a/src/flux_1_schnell/models/transformer/transformer.py
+++ b/src/flux_1_schnell/models/transformer/transformer.py
@@ -41,7 +41,7 @@ class Transformer(nn.Module):
         text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance)
         encoder_hidden_states = self.context_embedder(prompt_embeds)
         txt_ids = Transformer._prepare_text_ids(seq_len = prompt_embeds.shape[1])
-        img_ids = Transformer._prepare_latent_image_ids(config.width, config.height)
+        img_ids = Transformer._prepare_latent_image_ids(config.height, config.width)
         ids = mx.concatenate((txt_ids, img_ids), axis=1)
         image_rotary_emb = self.pos_embed.forward(ids)
 
@@ -69,14 +69,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