From 65c479b0919ae6c9d11c7e86c1ad2b5234ea3e3c Mon Sep 17 00:00:00 2001 From: Fabio Date: Sat, 17 Aug 2024 22:39:52 +0200 Subject: [PATCH 1/2] 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) From e76bb3d2849905d21363a5a9257d5430e265854c Mon Sep 17 00:00:00 2001 From: Fabio Date: Sat, 17 Aug 2024 22:50:35 +0200 Subject: [PATCH 2/2] change naming --- .../latent_creator/latent_creator.py | 6 +++--- .../models/transformer/transformer.py | 14 +++++++------- 2 files changed, 10 insertions(+), 10 deletions(-) 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 e0d9434..7676885 100644 --- a/src/flux_1_schnell/models/transformer/transformer.py +++ b/src/flux_1_schnell/models/transformer/transformer.py @@ -67,14 +67,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