From 68cbaab81645a5fc772be61308e5e77472e6ced1 Mon Sep 17 00:00:00 2001 From: Fabio Date: Sun, 18 Aug 2024 23:31:54 +0200 Subject: [PATCH] fix type and add main_dev --- .gitignore | 4 ---- main.py | 6 +++--- src/flux_1_schnell/config/config.py | 3 ++- src/flux_1_schnell/flux.py | 7 +++---- src/flux_1_schnell/models/transformer/transformer.py | 4 ++-- 5 files changed, 10 insertions(+), 14 deletions(-) diff --git a/.gitignore b/.gitignore index e7b1ada..ee9ae73 100644 --- a/.gitignore +++ b/.gitignore @@ -11,7 +11,3 @@ *.png *.jpg *.pyc - - -*.pt -trial.py \ No newline at end of file diff --git a/main.py b/main.py index aa11963..6d24b9b 100644 --- a/main.py +++ b/main.py @@ -6,15 +6,15 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src'))) from flux_1_schnell.config.config import Config from flux_1_schnell.flux import Flux1 -flux = Flux1("black-forest-labs/FLUX.1-dev") +flux = Flux1("black-forest-labs/FLUX.1-dev", max_sequence_length=512) image = flux.generate_image( seed=3, 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=20, - width=256, - height=256, + height=768, + width=1360, guidance=3.5, ) ) diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py index 0335128..a88532e 100644 --- a/src/flux_1_schnell/config/config.py +++ b/src/flux_1_schnell/config/config.py @@ -5,7 +5,7 @@ import logging log = logging.getLogger(__name__) class Config: - precision: mx.Dtype = mx.float16 + precision: mx.Dtype = mx.bfloat16 num_train_steps = 1000 def __init__( @@ -37,3 +37,4 @@ class Config: b = y1 - m * x1 mu = m * self.width * self.height / 256 + b self.sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / self.sigmas - 1)) + self.sigmas[-1] = 0 diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py index 7a02bdc..39cbbea 100644 --- a/src/flux_1_schnell/flux.py +++ b/src/flux_1_schnell/flux.py @@ -18,11 +18,10 @@ from flux_1_schnell.weights.weight_handler import WeightHandler class Flux1: - def __init__(self, repo_id: str): + def __init__(self, repo_id: str, max_sequence_length: int = 512): self.is_dev = "FLUX.1-dev" in repo_id - max_t5_length = 512 if self.is_dev else 256 - tokenizers = TokenizerHandler.load_from_disk_or_huggingface(repo_id, max_t5_length) - self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=max_t5_length) + tokenizers = TokenizerHandler.load_from_disk_or_huggingface(repo_id, max_sequence_length) + self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=max_sequence_length) self.clip_tokenizer = TokenizerCLIP(tokenizers.clip) weights = WeightHandler.load_from_disk_or_huggingface(repo_id) diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py index fa4e356..f5ce147 100644 --- a/src/flux_1_schnell/models/transformer/transformer.py +++ b/src/flux_1_schnell/models/transformer/transformer.py @@ -35,9 +35,9 @@ class Transformer(nn.Module): config: Config ) -> mx.array: time_step = config.sigmas[t] * config.num_train_steps - time_step = mx.broadcast_to(time_step, (1,)) + time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) hidden_states = self.x_embedder(hidden_states) - guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)) + guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision) text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance) encoder_hidden_states = self.context_embedder(prompt_embeds) txt_ids = Transformer._prepare_text_ids(seq_len = prompt_embeds.shape[1])