fix type and add main_dev

.gitignore

@@ -11,7 +11,3 @@
 *.png
 *.jpg
 *.pyc
-
-
-*.pt
-trial.py

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

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

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)

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