diff --git a/README.md b/README.md
index 1d1d1f8..06bebbe 100644
--- a/README.md
+++ b/README.md
@@ -51,6 +51,7 @@ sys.path.append("/path/to/mflux/src")
 
 from flux_1_schnell.config.config import Config
 from flux_1_schnell.flux import Flux1Schnell
+from flux_1_schnell.post_processing.image_util import ImageUtil
 
 flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell")
 
@@ -62,7 +63,7 @@ image = flux.generate_image(
    )
 )
 
-image.save("image.png")
+ImageUtil.save_image(image, "image.png")
 ```
 
 If the model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell model).
diff --git a/main.py b/main.py
index 00567f2..5cf39bb 100644
--- a/main.py
+++ b/main.py
@@ -19,4 +19,4 @@ image = flux.generate_image(
     )
 )
 
-ImageUtil.save_image(image, "image.png")
\ No newline at end of file
+ImageUtil.save_image(image, "image.png")
diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py
index a88532e..3524b9c 100644
--- a/src/flux_1_schnell/config/config.py
+++ b/src/flux_1_schnell/config/config.py
@@ -4,6 +4,7 @@ import logging
 
 log = logging.getLogger(__name__)
 
+
 class Config:
     precision: mx.Dtype = mx.bfloat16
     num_train_steps = 1000
diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py
index 39cbbea..f5eb5dd 100644
--- a/src/flux_1_schnell/flux.py
+++ b/src/flux_1_schnell/flux.py
@@ -59,10 +59,10 @@ class Flux1:
         return ImageUtil.to_image(decoded)
 
     @staticmethod
-    def _unpack_latents(latents: mx.array, width: int, height: int) -> mx.array:
-        latents = mx.reshape(latents, (1, height//16, width//16, 16, 2, 2))
+    def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
+        latents = mx.reshape(latents, (1, width // 16, height // 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))
+        latents = mx.reshape(latents, (1, 16, width // 16 * 2, height // 16 * 2))
         return latents
 
     def encode(self, path: str) -> mx.array:
diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py
index f5ce147..6638501 100644
--- a/src/flux_1_schnell/models/transformer/transformer.py
+++ b/src/flux_1_schnell/models/transformer/transformer.py
@@ -69,14 +69,14 @@ class Transformer(nn.Module):
         return noise
 
     @staticmethod
-    def _prepare_latent_image_ids(width: int, height: int) -> mx.array:
-        latent_height = height // 16
+    def _prepare_latent_image_ids(height: int, width: int) -> mx.array:
         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_height = height // 16
+        latent_image_ids = mx.zeros((latent_width, latent_height, 3))
+        latent_image_ids = latent_image_ids.at[:, :, 1].add(mx.arange(0, latent_width)[:, None])
+        latent_image_ids = latent_image_ids.at[:, :, 2].add(mx.arange(0, latent_height)[None, :])
         latent_image_ids = mx.repeat(latent_image_ids[None, :], 1, axis=0)
-        latent_image_ids = mx.reshape(latent_image_ids, (1, latent_height * latent_width, 3))
+        latent_image_ids = mx.reshape(latent_image_ids, (1, latent_width * latent_height, 3))
         return latent_image_ids
 
     @staticmethod
diff --git a/src/flux_1_schnell/post_processing/image_util.py b/src/flux_1_schnell/post_processing/image_util.py
index 95f175a..9cc282a 100644
--- a/src/flux_1_schnell/post_processing/image_util.py
+++ b/src/flux_1_schnell/post_processing/image_util.py
@@ -1,8 +1,13 @@
+import logging
+from pathlib import Path
+
 import PIL
 import mlx.core as mx
 import numpy as np
 from PIL import Image
 
+log = logging.getLogger(__name__)
+
 
 class ImageUtil:
 
@@ -53,3 +58,23 @@ class ImageUtil:
     def resize(image):
         image = image.resize((1024, 1024), resample=PIL.Image.LANCZOS)
         return image
+
+    @staticmethod
+    def save_image(image: Image.Image, path: str) -> None:
+        file_path = Path(path)
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_name = file_path.stem
+        file_extension = file_path.suffix
+
+        # If a file already exists, create a new name with a counter
+        counter = 1
+        while file_path.exists():
+            new_name = f"{file_name}({counter}){file_extension}"
+            file_path = file_path.with_name(new_name)
+            counter += 1
+
+        try:
+            image.save(file_path)
+            log.info(f"Image saved successfully at: {file_path}")
+        except Exception as e:
+            log.info(f"Error saving image: {e}")