From 766162283935588157b1495fbc50c7d33ad653e6 Mon Sep 17 00:00:00 2001 From: Fabio Peruzzo Date: Fri, 16 Aug 2024 16:45:49 +0200 Subject: [PATCH 1/3] allow different image resolutions --- README.md | 1 - main.py | 2 ++ src/flux_1_schnell/config/config.py | 8 ++++++++ src/flux_1_schnell/flux.py | 10 +++++----- .../latent_creator/latent_creator.py | 6 ++++-- .../models/transformer/transformer.py | 14 ++++++++------ 6 files changed, 27 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index c3f34fc..1d1d1f8 100644 --- a/README.md +++ b/README.md @@ -129,7 +129,6 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots - Images are generated one by one. - Negative prompts not supported. -- Image resolution is (1024, 1024) ### TODO diff --git a/main.py b/main.py index 7d6385b..60273f3 100644 --- a/main.py +++ b/main.py @@ -13,6 +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, ) ) diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py index fd2039d..a9ac241 100644 --- a/src/flux_1_schnell/config/config.py +++ b/src/flux_1_schnell/config/config.py @@ -1,6 +1,8 @@ import mlx.core as mx import numpy as np +import logging +log = logging.getLogger(__name__) class Config: precision: mx.Dtype = mx.float16 @@ -9,7 +11,13 @@ class Config: def __init__( self, num_inference_steps: int = 4, + width: int = 1024, + height: int = 1024, ): + 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) base_sigmas = Config.base_sigmas(num_inference_steps) self.num_inference_steps = num_inference_steps self.time_steps = base_sigmas * self.num_train_steps diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py index 5d3b367..d2d3299 100644 --- a/src/flux_1_schnell/flux.py +++ b/src/flux_1_schnell/flux.py @@ -31,7 +31,7 @@ class Flux1Schnell: self.clip_text_encoder = CLIPEncoder(weights.clip_encoder) def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image: - latents = LatentCreator.create(seed) + latents = LatentCreator.create(config.height, config.width, seed) t5_tokens = self.t5_tokenizer.tokenize(prompt) clip_tokens = self.clip_tokenizer.tokenize(prompt) @@ -56,15 +56,15 @@ class Flux1Schnell: mx.eval(latents) - latents = Flux1Schnell._unpack_latents(latents) + latents = Flux1Schnell._unpack_latents(latents, config.width, config.height) decoded = self.vae.decode(latents) return ImageUtil.to_image(decoded) @staticmethod - def _unpack_latents(latents): - latents = mx.reshape(latents, (1, 64, 64, 16, 2, 2)) + def _unpack_latents(latents, width, height): + 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, 128, 128)) + latents = mx.reshape(latents, (1, 16, height//16 *2, width//16 * 2)) return latents def encode(self, path: str) -> mx.array: diff --git a/src/flux_1_schnell/latent_creator/latent_creator.py b/src/flux_1_schnell/latent_creator/latent_creator.py index c4d3b1d..82309ee 100644 --- a/src/flux_1_schnell/latent_creator/latent_creator.py +++ b/src/flux_1_schnell/latent_creator/latent_creator.py @@ -4,6 +4,8 @@ import mlx.core as mx class LatentCreator: @staticmethod - def create(seed: int) -> (mx.array, mx.array): - latents = mx.random.normal(shape=[1, 4096, 64], key=mx.random.key(seed)) + 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)) return latents diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py index 94ca0ce..4bb779a 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() + img_ids = Transformer._prepare_latent_image_ids(config.width, config.height) ids = mx.concatenate((txt_ids, img_ids), axis=1) image_rotary_emb = self.pos_embed.forward(ids) @@ -67,12 +67,14 @@ class Transformer(nn.Module): return noise @staticmethod - def _prepare_latent_image_ids() -> mx.array: - latent_image_ids = mx.zeros((64, 64, 3)) - latent_image_ids = latent_image_ids.at[:, :, 1].add(mx.arange(0, 64)[:, None]) - latent_image_ids = latent_image_ids.at[:, :, 2].add(mx.arange(0, 64)[None, :]) + 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, :]) latent_image_ids = mx.repeat(latent_image_ids[None, :], 1, axis=0) - latent_image_ids = mx.reshape(latent_image_ids, (1, 4096, 3)) + latent_image_ids = mx.reshape(latent_image_ids, (1, r_h * r_w, 3)) return latent_image_ids @staticmethod From 65c479b0919ae6c9d11c7e86c1ad2b5234ea3e3c Mon Sep 17 00:00:00 2001 From: Fabio Date: Sat, 17 Aug 2024 22:39:52 +0200 Subject: [PATCH 2/3] 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 3/3] 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