fix height <-> width and other minor corrections

.gitignore

@@ -10,3 +10,4 @@
 .venv
 *.png
 *.jpg
+*.pyc

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,
     )
 )
 

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)

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))

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)