Merge pull request #31 from Xuzzo/fix/width_height_mixup

Fix width and height in config

src/flux_1/config/config.py

@@ -17,7 +17,7 @@ class Config:
     ):
         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)
+        self.width = 16 * (width // 16)
+        self.height = 16 * (height // 16)
         self.num_inference_steps = num_inference_steps
         self.guidance = guidance

src/flux_1/flux.py

@@ -104,9 +104,9 @@ class Flux1:
 
     @staticmethod
     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.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, width // 16 * 2, height // 16 * 2))
+        latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
         return latents
 
     @staticmethod

src/flux_1/models/transformer/transformer.py

@@ -69,9 +69,9 @@ class Transformer(nn.Module):
     def _prepare_latent_image_ids(height: int, width: int) -> mx.array:
         latent_width = width // 16
         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.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, latent_width * latent_height, 3))
         return latent_image_ids