From 65c479b0919ae6c9d11c7e86c1ad2b5234ea3e3c Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Sat, 17 Aug 2024 22:39:52 +0200
Subject: [PATCH] fix height <-> width and other minor corrections

---
 .gitignore                                           | 1 +
 main.py                                              | 4 ++--
 src/flux_1_schnell/config/config.py                  | 2 +-
 src/flux_1_schnell/flux.py                           | 4 ++--
 src/flux_1_schnell/models/transformer/transformer.py | 2 +-
 5 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/.gitignore b/.gitignore
index fd92c93..ee9ae73 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,3 +10,4 @@
 .venv
 *.png
 *.jpg
+*.pyc
diff --git a/main.py b/main.py
index 60273f3..ba639a1 100644
--- a/main.py
+++ b/main.py
@@ -13,8 +13,8 @@ image = flux.generate_image(
     prompt="Luxury food photograph of a birthday cake. In the middle it has three candles shaped like letters spelling the word 'MLX'. It has perfect lighting and a cozy background with big bokeh and shallow depth of field. The mood is a sunset balcony in tuscany. The photo is taken from the side of the cake. The scene is complemented by a warm, inviting light that highlights the textures and colors of the ingredients, giving it an appetizing and elegant look.",
     config=Config(
         num_inference_steps=2,
-        width=256,
-        height=256,
+        height=768,
+        width=1360,
     )
 )
 
diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py
index a9ac241..2a01fd2 100644
--- a/src/flux_1_schnell/config/config.py
+++ b/src/flux_1_schnell/config/config.py
@@ -14,7 +14,7 @@ class Config:
             width: int = 1024,
             height: int = 1024,
     ):
-        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 d2d3299..50dade0 100644
--- a/src/flux_1_schnell/flux.py
+++ b/src/flux_1_schnell/flux.py
@@ -56,12 +56,12 @@ class Flux1Schnell:
 
             mx.eval(latents)
 
-        latents = Flux1Schnell._unpack_latents(latents, config.width, config.height)
+        latents = Flux1Schnell._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/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py
index 4bb779a..e0d9434 100644
--- a/src/flux_1_schnell/models/transformer/transformer.py
+++ b/src/flux_1_schnell/models/transformer/transformer.py
@@ -39,7 +39,7 @@ class Transformer(nn.Module):
         text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds)
         encoder_hidden_states = self.context_embedder(prompt_embeds)
         txt_ids = Transformer._prepare_text_ids()
-        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)