Initial release

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,27 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPep8Inspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="ignoredErrors">
+        <list>
+          <option value="E501" />
+        </list>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="ignoredErrors">
+        <list>
+          <option value="N806" />
+        </list>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyUnresolvedReferencesInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredIdentifiers">
+        <list>
+          <option value="list.__await__" />
+        </list>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyUnusedLocalInspection" enabled="false" level="WEAK WARNING" enabled_by_default="false" />
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/mflux.iml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.10 (mlx-examples)" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/mflux.iml" filepath="$PROJECT_DIR$/.idea/mflux.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

README.md

@@ -1 +1,119 @@
-# mflux
+
+![image](src/flux_1_schnell/assets/logo.png)
+*A MLX port of FLUX based on the Huggingface Diffusers implementation.*
+
+### About
+
+Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black Forest Labs](https://blackforestlabs.ai) locally on your Mac!
+
+### Philosophy
+
+MFLUX (MacFLUX) is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx). 
+MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used
+except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models 
+(thereby avoiding too many abstractions). MFLUX priorities readability over generality and performance.
+
+All models are implemented from scratch in MLX and only the tokenizers are used via the 
+[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies 
+like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) for simple image post-processing.
+
+### Models
+
+- [x] FLUX.1-Scnhell
+- [ ] FLUX.1-Dev
+
+### Prerequisites
+
+- The required libraries are listed in [requirements.txt](requirements.txt). 
+- Development and testing was done with Python 3.12
+- Download the weights and tokenizers at [Huggingface/black-forest-labs](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main).
+    Point to the location of the root folder as shown in the example below:
+
+### Generating an image
+
+Run [main.py](/src/flux_1_schnell/main.py) or make a new script like the following
+(make sure to correctly define the root path to wherever the model is saved and where you want the images to be stored).
+
+```python
+from flux_1_schnell.config.config import Config
+from flux_1_schnell.models.flux import Flux1Schnell
+
+flux = Flux1Schnell("/Users/filipstrand/.cache/FLUX.1-schnell/")
+
+image = flux.generate_image(
+    seed=42,
+    prompt="Luxury food photograph",
+    config=Config(
+        num_inference_steps=2,
+    )
+)
+
+image.save(f"/Users/filipstrand/Desktop/image.png")
+```
+
+Generating a single image (with 2 inference steps, Schnell model) takes between 2 and 3 minutes. This implementation has been tested on two Macbook Pro machines: 
+- 2021 M1 Pro (32 GB) 
+- 2023 M2 Max (32 GB)
+
+### Equivalent to Diffusers implementation 
+
+There is only a single source of randomness when generating an image: The initial latent array. 
+In this implementation, this initial latent is fully deterministically controlled by the input `seed` parameter. 
+However, if we were to import a fixed instance of this latent array saved from the Diffusers implementation, then MFLUX will produce an identical image to the Diffusers implementation (assuming a fixed prompt and using the default parameter settings in the Diffusers setup).
+
+
+The images-pairs below illustrate this equivalence. 
+In these both cases the Schnell model was run for 2 time steps, 
+the Diffusers implementation ran in CPU mode. 
+The precision for MFLUX can be set in the [Config](src/flux_1_schnell/config/config.py) class. 
+There is typically a noticeable but very small difference in the final image when switching between 16bit and 32bit precision.
+
+---
+```
+Luxury food photograph
+```
+![image](src/flux_1_schnell/assets/comparison1.jpg)
+
+---
+```
+detailed cinematic dof render of an old dusty detailed CRT monitor on a wooden desk in a dim room with items around, messy dirty room. On the screen are the letters "FLUX" glowing softly. High detail hard surface render
+```
+![image](src/flux_1_schnell/assets/comparison2.jpg)
+
+---
+
+```
+photorealistic, lotr, A tiny red dragon curled up asleep inside a nest, (Soft Focus) , (f_stop 2.8) , (focal_length 50mm) macro lens f/2. 8, medieval wizard table, (pastel) colors, (cozy) morning light filtering through a nearby window, (whimsical) steam shapes, captured with a (Canon EOS R5) , highlighting (serene) comfort, medieval, dnd, rpg, 3d, 16K, 8K
+```
+![image](src/flux_1_schnell/assets/comparison3.jpg)
+
+---
+
+
+```
+A weathered fisherman in his early 60s stands on the deck of his boat, gazing out at a stormy sea. He has a thick, salt-and-pepper beard, deep-set blue eyes, and skin tanned and creased from years of sun exposure. He's wearing a yellow raincoat and hat, with water droplets clinging to the fabric. Behind him, dark clouds loom ominously, and waves crash against the side of the boat. The overall atmosphere is one of tension and respect for the power of nature.
+```
+![image](src/flux_1_schnell/assets/comparison4.jpg)
+
+---
+
+```
+Luxury food photograph of an italian Linguine pasta alle vongole dish with lots of clams. 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 plate. The pasta is shiny with sprinkled parmesan cheese and basil leaves on top. 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.
+```
+![image](src/flux_1_schnell/assets/comparison5.jpg)
+
+---
+
+
+
+### Current limitations
+
+- Images are generated one by one.
+- Negative prompts not supported.
+- Image resolution is (1024, 1024)
+
+### TODO
+
+- FLUX Dev implementation
+- LoRA adapters
+- Command line args

requirements.txt

@@ -0,0 +1,6 @@
+mlx>=0.16.0
+numpy>=2.0.1
+pillow>=10.4.0
+transformers>=4.44.0
+sentencepiece>=0.2.0
+torch>=2.3.1

src/flux_1_schnell/config/config.py

@@ -0,0 +1,22 @@
+import mlx.core as mx
+import numpy as np
+
+
+class Config:
+    precision: mx.Dtype = mx.float16
+    num_train_steps = 1000
+
+    def __init__(
+            self,
+            num_inference_steps: int = 4,
+    ):
+        base_sigmas = Config.base_sigmas(num_inference_steps)
+        self.num_inference_steps = num_inference_steps
+        self.time_steps = base_sigmas * self.num_train_steps
+        self.sigmas = mx.concatenate([base_sigmas, mx.zeros(1)])
+
+    @staticmethod
+    def base_sigmas(num_inference_steps):
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        sigmas = mx.array(sigmas).astype(mx.float32)
+        return sigmas

src/flux_1_schnell/latent_creator/latent_creator.py

@@ -0,0 +1,9 @@
+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))
+        return latents

src/flux_1_schnell/main.py

@@ -0,0 +1,15 @@
+from flux_1_schnell.config.config import Config
+from flux_1_schnell.models.flux import Flux1Schnell
+
+flux = Flux1Schnell("/Users/filipstrand/.cache/FLUX.1-schnell/")
+
+for i in range(0, 10):
+    image = flux.generate_image(
+        seed=i,
+        prompt="Luxury food photograph of an italian Linguine pasta alle vongole dish with lots of clams. 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 plate. The pasta is shiny with sprinkled parmesan cheese and basil leaves on top. 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,
+        )
+    )
+
+    image.save(f"/Users/filipstrand/Desktop/image_{i}.png")

src/flux_1_schnell/models/flux.py

@@ -0,0 +1,77 @@
+import PIL
+import mlx.core as mx
+from PIL import Image
+
+from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
+from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
+
+from flux_1_schnell.config.config import Config
+from flux_1_schnell.latent_creator.latent_creator import LatentCreator
+from flux_1_schnell.models.text_encoder.text_encoder import TextEncoder
+from flux_1_schnell.models.transformer.transformer import Transformer
+from flux_1_schnell.models.vae.vae import VAE
+from flux_1_schnell.post_processing.image_util import ImageUtil
+from flux_1_schnell.scheduler.scheduler import FlowMatchEulerDiscreteNoiseScheduler
+from flux_1_schnell.tokenizer.tokenizer_handler import TokenizerHandler
+from flux_1_schnell.weights.weight_handler import WeightHandler
+
+
+class Flux1Schnell:
+
+    def __init__(self, root_path: str):
+        tokenizers = TokenizerHandler.load_from_disk_via_huggingface_transformers(root_path)
+        self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
+        self.t5_tokenizer = TokenizerT5(tokenizers.t5)
+
+        weights = WeightHandler.load_from_disk(root_path)
+        self.vae = VAE(weights.vae)
+        self.transformer = Transformer(weights.transformer)
+        self.clip_text_encoder = CLIPEncoder(weights.clip_encoder)
+        self.t5_text_encoder = T5Encoder(weights.t5_encoder)
+
+    def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image:
+        latents = LatentCreator.create(seed)
+        prompt_embeds, pooled_prompt_embeds = TextEncoder.encode(
+            prompt=prompt,
+            clip_tokenizer=self.clip_tokenizer,
+            t5_tokenizer=self.t5_tokenizer,
+            clip_text_encoder=self.clip_text_encoder,
+            t5_text_encoder=self.t5_text_encoder
+        )
+
+        for t in range(config.num_inference_steps):
+            noise = self.transformer.predict(
+                t=t,
+                prompt_embeds=prompt_embeds,
+                pooled_prompt_embeds=pooled_prompt_embeds,
+                hidden_states=latents,
+                config=config
+            )
+
+            latents = FlowMatchEulerDiscreteNoiseScheduler.denoise(
+                t=t,
+                noise=noise,
+                latent=latents,
+                config=config
+            )
+
+        latents = Flux1Schnell._unpack_latents(latents)
+        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))
+        latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
+        latents = mx.reshape(latents, (1, 16, 128, 128))
+        return latents
+
+    def encode(self, path: str) -> mx.array:
+        array = ImageUtil.to_array(Image.open(path))
+        return self.vae.encode(array)
+
+    def decode(self, code: mx.array) -> PIL.Image.Image:
+        decoded = self.vae.decode(code)
+        return ImageUtil.to_image(decoded)

src/flux_1_schnell/models/text_encoder/clip_encoder/clip_embeddings.py

@@ -0,0 +1,20 @@
+from mlx import nn
+import mlx.core as mx
+
+from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
+
+
+class CLIPEmbeddings(nn.Module):
+
+    def __init__(self, dims: int):
+        super().__init__()
+        self.position_embedding = nn.Embedding(num_embeddings=TokenizerCLIP.MAX_TOKEN_LENGTH, dims=dims)
+        self.token_embedding = nn.Embedding(num_embeddings=49408, dims=dims)
+
+    def forward(self, tokens: mx.array) -> mx.array:
+        seq_length = tokens.shape[-1]
+        position_ids = mx.arange(0, seq_length).reshape(1, seq_length)
+        inputs_embeds = self.token_embedding(tokens)
+        position_embeddings = self.position_embedding(position_ids)
+        embeddings = inputs_embeds + position_embeddings
+        return embeddings

src/flux_1_schnell/models/text_encoder/clip_encoder/clip_encoder.py

@@ -0,0 +1,18 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.text_encoder.clip_encoder.clip_text_model import CLIPTextModel
+
+
+class CLIPEncoder(nn.Module):
+
+    def __init__(self, weights: dict):
+        super().__init__()
+        self.text_model = CLIPTextModel(dims=768, num_encoder_layers=12)
+
+        # Load the weights after all components are initialized
+        self.update(weights)
+
+    def forward(self, tokens: mx.array) -> mx.array:
+        pooled_output = self.text_model.forward(tokens)
+        return pooled_output

src/flux_1_schnell/models/text_encoder/clip_encoder/clip_encoder_layer.py

@@ -0,0 +1,26 @@
+from mlx import nn
+import mlx.core as mx
+
+from flux_1_schnell.models.text_encoder.clip_encoder.clip_mlp import CLIPMLP
+from flux_1_schnell.models.text_encoder.clip_encoder.clip_sdpa_attention import CLIPSdpaAttention
+
+
+class CLIPEncoderLayer(nn.Module):
+
+    def __init__(self, layer: int):
+        super().__init__()
+        self.self_attn = CLIPSdpaAttention()
+        self.layer_norm1 = nn.LayerNorm(dims=768)
+        self.mlp = CLIPMLP()
+        self.layer_norm2 = nn.LayerNorm(dims=768)
+
+    def forward(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array:
+        residual = hidden_states
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states = self.self_attn.forward(hidden_states, causal_attention_mask)
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp.forward(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states

src/flux_1_schnell/models/text_encoder/clip_encoder/clip_mlp.py

@@ -0,0 +1,20 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class CLIPMLP(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dims=768, output_dims=3072)
+        self.fc2 = nn.Linear(input_dims=3072, output_dims=768)
+
+    def forward(self, hidden_states: mx.array) -> mx.array:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = CLIPMLP.quick_gelu(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+    @staticmethod
+    def quick_gelu(input_array: mx.array) -> mx.array:
+        return input_array * nn.sigmoid(1.702 * input_array)

src/flux_1_schnell/models/text_encoder/clip_encoder/clip_sdpa_attention.py

@@ -0,0 +1,44 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class CLIPSdpaAttention(nn.Module):
+    head_dimension = 64
+    batch_size = 1
+    num_heads = 12
+
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(input_dims=768, output_dims=768)
+        self.k_proj = nn.Linear(input_dims=768, output_dims=768)
+        self.v_proj = nn.Linear(input_dims=768, output_dims=768)
+        self.out_proj = nn.Linear(input_dims=768, output_dims=768)
+
+    def forward(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array:
+        query = self.q_proj(hidden_states)
+        key = self.k_proj(hidden_states)
+        value = self.v_proj(hidden_states)
+
+        query = CLIPSdpaAttention.reshape_and_transpose(query, self.batch_size, self.num_heads, self.head_dimension)
+        key = CLIPSdpaAttention.reshape_and_transpose(key, self.batch_size, self.num_heads, self.head_dimension)
+        value = CLIPSdpaAttention.reshape_and_transpose(value, self.batch_size, self.num_heads, self.head_dimension)
+
+        hidden_states = CLIPSdpaAttention.masked_attention(query, key, value, causal_attention_mask)
+        hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3))
+        hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension))
+
+        hidden_states = self.out_proj(hidden_states)
+        return hidden_states
+
+    @staticmethod
+    def masked_attention(query, key, value, mask):
+        scale = 1 / mx.sqrt(query.shape[-1])
+        scores = (query * scale) @ key.transpose(0, 1, 3, 2)
+        scores = scores + mask
+        attn = mx.softmax(scores, axis=-1)
+        hidden_states = (attn @ value)
+        return hidden_states
+
+    @staticmethod
+    def reshape_and_transpose(x, batch_size, num_heads, head_dim):
+        return mx.transpose(mx.reshape(x, (batch_size, -1, num_heads, head_dim)), (0, 2, 1, 3))

src/flux_1_schnell/models/text_encoder/clip_encoder/clip_text_model.py

@@ -0,0 +1,33 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.text_encoder.clip_encoder.clip_embeddings import CLIPEmbeddings
+from flux_1_schnell.models.text_encoder.clip_encoder.encoder_clip import EncoderCLIP
+
+
+class CLIPTextModel(nn.Module):
+
+    def __init__(self, dims: int, num_encoder_layers: int):
+        super().__init__()
+        self.encoder = EncoderCLIP(num_encoder_layers)
+        self.embeddings = CLIPEmbeddings(dims)
+        self.final_layer_norm = nn.LayerNorm(dims=768)
+
+    def forward(self, tokens: mx.array) -> (mx.array, mx.array):
+        hidden_states = self.embeddings.forward(tokens)
+        causal_attention_mask = CLIPTextModel.create_causal_attention_mask(hidden_states.shape)
+        encoder_outputs = self.encoder.forward(hidden_states, causal_attention_mask)
+        last_hidden_state = self.final_layer_norm(encoder_outputs)
+        pooled_output = last_hidden_state[0, mx.argmax(tokens, axis=-1)]
+        return pooled_output
+
+    @staticmethod
+    def create_causal_attention_mask(input_shape: tuple) -> mx.array:
+        batch_size, query_length, _ = input_shape
+        key_value_length = query_length
+        mask = mx.tril(x=mx.ones((query_length, key_value_length)), k=0)
+        mask = 1 - mask
+        mask = mask * -3.4e38
+        mask = mask.reshape((1, 1, query_length, key_value_length))
+        mask = mx.broadcast_to(mask, (batch_size, 1, query_length, key_value_length))
+        return mask

src/flux_1_schnell/models/text_encoder/clip_encoder/encoder_clip.py

@@ -0,0 +1,20 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder_layer import CLIPEncoderLayer
+
+
+class EncoderCLIP(nn.Module):
+
+    def __init__(self, num_encoder_layers: int):
+        super().__init__()
+        self.layers = [CLIPEncoderLayer(i) for i in range(num_encoder_layers)]
+
+    def forward(self, tokens: mx.array, causal_attention_mask: mx.array) -> mx.array:
+        hidden_states = tokens
+        for encoder_layer in self.layers:
+            hidden_states = encoder_layer.forward(
+                hidden_states,
+                causal_attention_mask
+            )
+        return hidden_states

src/flux_1_schnell/models/text_encoder/t5_encoder/t5_attention.py

@@ -0,0 +1,19 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_self_attention import T5SelfAttention
+
+
+class T5Attention(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.SelfAttention = T5SelfAttention()
+        self.layer_norm = T5LayerNorm()
+
+    def forward(self, hidden_states: mx.array) -> mx.array:
+        normed_hidden_states = self.layer_norm.forward(hidden_states)
+        attention_output = self.SelfAttention.forward(normed_hidden_states)
+        hidden_states = hidden_states + attention_output
+        return hidden_states

src/flux_1_schnell/models/text_encoder/t5_encoder/t5_block.py

@@ -0,0 +1,19 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_attention import T5Attention
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_feed_forward import T5FeedForward
+
+
+class T5Block(nn.Module):
+
+    def __init__(self, layer: int):
+        super().__init__()
+        self.attention = T5Attention()
+        self.ff = T5FeedForward()
+
+    def forward(self, hidden_states: mx.array) -> mx.array:
+        hidden_states = self.attention.forward(hidden_states)
+        hidden_states = self.ff.forward(hidden_states)
+        outputs = hidden_states
+        return outputs

src/flux_1_schnell/models/text_encoder/t5_encoder/t5_dense_relu_dense.py

@@ -0,0 +1,24 @@
+import math
+
+from mlx import nn
+import mlx.core as mx
+
+
+class T5DenseReluDense(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.wi_0 = nn.Linear(4096, 10240, bias=False)
+        self.wi_1 = nn.Linear(4096, 10240, bias=False)
+        self.wo = nn.Linear(10240, 4096, bias=False)
+
+    def forward(self, hidden_states: mx.array) -> mx.array:
+        hidden_gelu = T5DenseReluDense.new_gelu(self.wi_0(hidden_states))
+        hidden_linear = self.wi_1(hidden_states)
+        hidden_states = hidden_gelu * hidden_linear
+        hidden_states = self.wo(hidden_states)
+        return hidden_states
+
+    @staticmethod
+    def new_gelu(input_array: mx.array) -> mx.array:
+        return 0.5 * input_array * (1.0 + mx.tanh(math.sqrt(2.0 / math.pi) * (input_array + 0.044715 * mx.power(input_array, 3.0))))

src/flux_1_schnell/models/text_encoder/t5_encoder/t5_encoder.py

@@ -0,0 +1,24 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_block import T5Block
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
+
+
+class T5Encoder(nn.Module):
+
+    def __init__(self, weights: dict):
+        super().__init__()
+        self.shared = nn.Embedding(num_embeddings=32128, dims=4096)
+        self.t5_blocks = [T5Block(i) for i in range(24)]
+        self.final_layer_norm = T5LayerNorm()
+
+        # Load the weights after all components are initialized
+        self.update(weights)
+
+    def forward(self, tokens: mx.array):
+        hidden_states = self.shared(tokens)
+        for block in self.t5_blocks:
+            hidden_states = block.forward(hidden_states)
+        hidden_states = self.final_layer_norm.forward(hidden_states)
+        return hidden_states

src/flux_1_schnell/models/text_encoder/t5_encoder/t5_feed_forward.py

@@ -0,0 +1,20 @@
+import math
+
+from mlx import nn
+import mlx.core as mx
+
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_dense_relu_dense import T5DenseReluDense
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
+
+
+class T5FeedForward(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.layer_norm = T5LayerNorm()
+        self.DenseReluDense = T5DenseReluDense()
+
+    def forward(self, hidden_states: mx.array) -> mx.array:
+        forwarded_states = self.layer_norm.forward(hidden_states)
+        forwarded_states = self.DenseReluDense.forward(forwarded_states)
+        return hidden_states + forwarded_states

src/flux_1_schnell/models/text_encoder/t5_encoder/t5_layer_norm.py

@@ -0,0 +1,15 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class T5LayerNorm(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.weight = mx.ones((4096,))
+        self.variance_epsilon = 1e-06
+
+    def forward(self, hidden_states: mx.array) -> mx.array:
+        variance = mx.mean(mx.power(hidden_states.astype(mx.float32), 2), axis=-1, keepdims=True)
+        hidden_states = hidden_states * mx.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states

src/flux_1_schnell/models/text_encoder/t5_encoder/t5_self_attention.py

@@ -0,0 +1,70 @@
+import math
+
+import mlx.core as mx
+from mlx import nn
+
+
+class T5SelfAttention(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.q = nn.Linear(4096, 4096, bias=False)
+        self.k = nn.Linear(4096, 4096, bias=False)
+        self.v = nn.Linear(4096, 4096, bias=False)
+        self.relative_attention_bias = nn.Embedding(32, 64)
+        self.o = nn.Linear(4096, 4096, bias=False)
+
+    def forward(self, hidden_states: mx.array) -> mx.array:
+        query_states = T5SelfAttention.shape(self.q(hidden_states))
+        key_states = T5SelfAttention.shape(self.k(hidden_states))
+        value_states = T5SelfAttention.shape(self.v(hidden_states))
+        scores = mx.matmul(query_states, mx.transpose(key_states, (0, 1, 3, 2)))
+        position_bias = self._compute_bias()
+        scores += position_bias
+        attn_weights = nn.softmax(scores, axis=-1)
+        attn_output = T5SelfAttention.un_shape(mx.matmul(attn_weights, value_states))
+        attn_output = self.o(attn_output)
+        return attn_output
+
+    @staticmethod
+    def shape(states):
+        return mx.transpose(mx.reshape(states, (1, -1, 64, 64)), (0, 2, 1, 3))
+
+    @staticmethod
+    def un_shape(states):
+        return mx.reshape(mx.transpose(states, (0, 2, 1, 3)), (1, -1, 4096))
+
+    def _compute_bias(self):
+        context_position = mx.arange(start=0, stop=256, step=1)[:, None]
+        memory_position = mx.arange(start=0, stop=256, step=1)[None, :]
+        relative_position = memory_position - context_position
+        relative_position_bucket = T5SelfAttention._relative_position_bucket(relative_position)
+        values = self.relative_attention_bias(relative_position_bucket)
+        values = mx.transpose(values, (2, 0, 1))
+        values = mx.expand_dims(values, 0)
+        return values
+
+    @staticmethod
+    def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
+        relative_buckets = mx.zeros_like(relative_position)
+        num_buckets //= 2
+        relative_buckets += mx.where(relative_position > 0, num_buckets, 0)
+        relative_position = mx.abs(relative_position)
+
+        # half of the buckets are for exact increments in positions
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+
+        # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
+        relative_position_if_large = max_exact + mx.floor(
+            mx.log(relative_position.astype(mx.float32) / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).astype(mx.int32)
+        relative_position_if_large = mx.minimum(
+            relative_position_if_large,
+            mx.full(relative_position_if_large.shape, num_buckets - 1)
+        )
+
+        relative_buckets += mx.where(is_small, relative_position, relative_position_if_large)
+        return relative_buckets

src/flux_1_schnell/models/text_encoder/text_encoder.py

@@ -0,0 +1,25 @@
+import mlx.core as mx
+
+from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from flux_1_schnell.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
+from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
+from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
+
+
+class TextEncoder:
+
+    @staticmethod
+    def encode(
+            prompt: str,
+            clip_tokenizer: TokenizerCLIP,
+            t5_tokenizer: TokenizerT5,
+            clip_text_encoder: CLIPEncoder,
+            t5_text_encoder: T5Encoder
+    ) -> (mx.array, mx.array):
+        clip_tokens = clip_tokenizer.tokenize(prompt)
+        pooled_prompt_embeds = clip_text_encoder.forward(clip_tokens)
+
+        t5_tokens = t5_tokenizer.tokenize(prompt)
+        prompt_embeds = t5_text_encoder.forward(t5_tokens)
+
+        return prompt_embeds, pooled_prompt_embeds

src/flux_1_schnell/models/transformer/ada_layer_norm_continous.py

@@ -0,0 +1,22 @@
+from mlx import nn
+import mlx.core as mx
+
+from flux_1_schnell.config.config import Config
+
+
+class AdaLayerNormContinuous(nn.Module):
+
+    def __init__(self, embedding_dim: int, conditioning_embedding_dim: int):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2)
+        self.norm = nn.LayerNorm(dims=embedding_dim, eps=1e-6, affine=False)
+
+    def forward(self, x: mx.array, text_embeddings: mx.array) -> mx.array:
+        text_embeddings = self.linear(nn.silu(text_embeddings).astype(Config.precision))
+        chunk_size = self.embedding_dim
+        scale = text_embeddings[:, 0*chunk_size:1*chunk_size]
+        shift = text_embeddings[:, 1*chunk_size:2*chunk_size]
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+

src/flux_1_schnell/models/transformer/ada_layer_norm_zero.py

@@ -0,0 +1,22 @@
+from mlx import nn
+import mlx.core as mx
+
+
+class AdaLayerNormZero(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(3072, 18432)
+        self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
+
+    def forward(self, x: mx.array, text_embeddings: mx.array):
+        text_embeddings = self.linear(nn.silu(text_embeddings))
+        chunk_size = 18432 // 6
+        shift_msa = text_embeddings[:, 0*chunk_size:1*chunk_size]
+        scale_msa = text_embeddings[:, 1*chunk_size:2*chunk_size]
+        gate_msa = text_embeddings[:, 2*chunk_size:3*chunk_size]
+        shift_mlp = text_embeddings[:, 3*chunk_size:4*chunk_size]
+        scale_mlp = text_embeddings[:, 4*chunk_size:5*chunk_size]
+        gate_mlp = text_embeddings[:, 5*chunk_size:6*chunk_size]
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp

src/flux_1_schnell/models/transformer/ada_layer_norm_zero_single.py

@@ -0,0 +1,19 @@
+from mlx import nn
+import mlx.core as mx
+
+
+class AdaLayerNormZeroSingle(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(3072, 3*3072)
+        self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
+
+    def forward(self, x: mx.array, text_embeddings: mx.array):
+        text_embeddings = self.linear(nn.silu(text_embeddings))
+        chunk_size = 9216 // 3
+        shift_msa = text_embeddings[:, 0*chunk_size:1*chunk_size]
+        scale_msa = text_embeddings[:, 1*chunk_size:2*chunk_size]
+        gate_msa = text_embeddings[:, 2*chunk_size:3*chunk_size]
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa

src/flux_1_schnell/models/transformer/embed_nd.py

@@ -0,0 +1,32 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class EmbedND(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.dim = 3072
+        self.theta = 10000
+        self.axes_dim = [16, 56, 56]
+
+    def forward(self, ids: mx.array) -> mx.array:
+        emb = mx.concatenate(
+            [EmbedND.rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(3)],
+            axis=-3,
+        )
+        return mx.expand_dims(emb, axis=1)
+
+    @staticmethod
+    def rope(pos: mx.array, dim: int, theta: float) -> mx.array:
+        scale = mx.arange(0, dim, 2, dtype=mx.float32) / dim
+        omega = 1.0 / (theta ** scale)
+        batch_size, seq_length = pos.shape
+        pos_expanded = mx.expand_dims(pos, axis=-1)
+        omega_expanded = mx.expand_dims(omega, axis=0)
+        out = pos_expanded * omega_expanded
+        cos_out = mx.cos(out)
+        sin_out = mx.sin(out)
+        stacked_out = mx.stack([cos_out, -sin_out, sin_out, cos_out], axis=-1)
+        out = mx.reshape(stacked_out, (batch_size, -1, dim // 2, 2, 2))
+        return out

src/flux_1_schnell/models/transformer/feed_forward.py

@@ -0,0 +1,17 @@
+from mlx import nn
+import mlx.core as mx
+
+
+class FeedForward(nn.Module):
+
+    def __init__(self, activation_function):
+        super().__init__()
+        self.linear1 = nn.Linear(3072, 6144)
+        self.linear2 = nn.Linear(6144, 3072)
+        self.activation_function = activation_function
+
+    def forward(self, hidden_states: mx.array) -> mx.array:
+        hidden_states = self.linear1(hidden_states)
+        hidden_states = self.activation_function(hidden_states)
+        hidden_states = self.linear2(hidden_states)
+        return hidden_states

src/flux_1_schnell/models/transformer/joint_attention.py

@@ -0,0 +1,88 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class JointAttention(nn.Module):
+    head_dimension = 128
+    batch_size = 1
+    num_heads = 24
+
+    def __init__(self):
+        super().__init__()
+        self.to_q = nn.Linear(3072, 3072)
+        self.to_k = nn.Linear(3072, 3072)
+        self.to_v = nn.Linear(3072, 3072)
+        self.to_out = [nn.Linear(3072, 3072)]
+        self.add_q_proj = nn.Linear(3072, 3072)
+        self.add_k_proj = nn.Linear(3072, 3072)
+        self.add_v_proj = nn.Linear(3072, 3072)
+        self.to_add_out = nn.Linear(3072, 3072)
+        self.norm_q = nn.RMSNorm(128)
+        self.norm_k = nn.RMSNorm(128)
+        self.norm_added_q = nn.RMSNorm(128)
+        self.norm_added_k = nn.RMSNorm(128)
+
+    def forward(
+            self,
+            hidden_states: mx.array,
+            encoder_hidden_states: mx.array,
+            image_rotary_emb: mx.array
+    ) -> (mx.array, mx.array):
+        residual = hidden_states
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        query = mx.transpose(mx.reshape(query, (1, -1, 24, 128)), (0, 2, 1, 3))
+        key = mx.transpose(mx.reshape(key, (1, -1, 24, 128)), (0, 2, 1, 3))
+        value = mx.transpose(mx.reshape(value, (1, -1, 24, 128)), (0, 2, 1, 3))
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        encoder_hidden_states_query_proj = self.add_q_proj(encoder_hidden_states)
+        encoder_hidden_states_key_proj = self.add_k_proj(encoder_hidden_states)
+        encoder_hidden_states_value_proj = self.add_v_proj(encoder_hidden_states)
+
+        encoder_hidden_states_query_proj = mx.transpose(mx.reshape(encoder_hidden_states_query_proj, (1, -1, 24, 128)), (0, 2, 1, 3))
+        encoder_hidden_states_key_proj = mx.transpose(mx.reshape(encoder_hidden_states_key_proj, (1, -1, 24, 128)), (0, 2, 1, 3))
+        encoder_hidden_states_value_proj = mx.transpose(mx.reshape(encoder_hidden_states_value_proj, (1, -1, 24, 128)),   (0, 2, 1, 3))
+
+        encoder_hidden_states_query_proj = self.norm_added_q(encoder_hidden_states_query_proj)
+        encoder_hidden_states_key_proj = self.norm_added_k(encoder_hidden_states_key_proj)
+
+        query = mx.concatenate([encoder_hidden_states_query_proj, query], axis=2)
+        key = mx.concatenate([encoder_hidden_states_key_proj, key], axis=2)
+        value = mx.concatenate([encoder_hidden_states_value_proj, value], axis=2)
+
+        query, key = JointAttention.apply_rope(query, key, image_rotary_emb)
+
+        hidden_states = JointAttention.attention(query, key, value)
+        hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3))
+        hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension))
+        encoder_hidden_states, hidden_states = (
+            hidden_states[:, : encoder_hidden_states.shape[1]],
+            hidden_states[:, encoder_hidden_states.shape[1]:],
+        )
+
+        hidden_states = self.to_out[0](hidden_states)
+        encoder_hidden_states = self.to_add_out(encoder_hidden_states)
+
+        return hidden_states, encoder_hidden_states
+
+    @staticmethod
+    def attention(query, key, value):
+        scale = 1 / mx.sqrt(query.shape[-1])
+        scores = (query * scale) @ key.transpose(0, 1, 3, 2)
+        attn = mx.softmax(scores, axis=-1)
+        hidden_states = (attn @ value)
+        return hidden_states
+
+    @staticmethod
+    def apply_rope(xq: mx.array, xk: mx.array, freqs_cis: mx.array):
+        xq_ = xq.astype(mx.float32).reshape(*xq.shape[:-1], -1, 1, 2)
+        xk_ = xk.astype(mx.float32).reshape(*xk.shape[:-1], -1, 1, 2)
+        xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+        xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+        return xq_out.reshape(*xq.shape).astype(mx.float32), xk_out.reshape(*xk.shape).astype(mx.float32)

src/flux_1_schnell/models/transformer/joint_transformer_block.py

@@ -0,0 +1,56 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.transformer.ada_layer_norm_zero import AdaLayerNormZero
+from flux_1_schnell.models.transformer.feed_forward import FeedForward
+from flux_1_schnell.models.transformer.joint_attention import JointAttention
+
+
+class JointTransformerBlock(nn.Module):
+
+    def __init__(self, layer):
+        super().__init__()
+        self.layer = layer
+        self.norm1 = AdaLayerNormZero()
+        self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
+        self.ff = FeedForward(activation_function=nn.gelu)
+        self.attn = JointAttention()
+        self.norm1_context = AdaLayerNormZero()
+        self.ff_context = FeedForward(activation_function=nn.gelu_approx)
+        self.norm2_context = nn.LayerNorm(dims=1536, eps=1e-6, affine=False)
+
+    def forward(
+            self,
+            hidden_states: mx.array,
+            encoder_hidden_states: mx.array,
+            text_embeddings: mx.array,
+            rotary_embeddings: mx.array
+    ) -> (mx.array, mx.array):
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1.forward(hidden_states, text_embeddings)
+
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context.forward(
+            x=encoder_hidden_states,
+            text_embeddings=text_embeddings
+        )
+
+        attn_output, context_attn_output = self.attn.forward(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=rotary_embeddings,
+        )
+
+        attn_output = mx.expand_dims(gate_msa, axis=1) * attn_output
+        hidden_states = hidden_states + attn_output
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff.forward(norm_hidden_states)
+        ff_output = mx.expand_dims(gate_mlp, axis=1) * ff_output
+        hidden_states = hidden_states + ff_output
+
+        context_attn_output = mx.expand_dims(c_gate_msa, axis=1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+        context_ff_output = self.ff_context.forward(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + mx.expand_dims(c_gate_mlp, axis=1) * context_ff_output
+        return encoder_hidden_states, hidden_states

src/flux_1_schnell/models/transformer/single_block_attention.py

@@ -0,0 +1,56 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class SingleBlockAttention(nn.Module):
+    head_dimension = 128
+    batch_size = 1
+    num_heads = 24
+
+    def __init__(self):
+        super().__init__()
+        self.to_q = nn.Linear(3072, 3072)
+        self.to_k = nn.Linear(3072, 3072)
+        self.to_v = nn.Linear(3072, 3072)
+        self.norm_q = nn.RMSNorm(128)
+        self.norm_k = nn.RMSNorm(128)
+
+    def forward(
+            self,
+            hidden_states: mx.array,
+            image_rotary_emb: mx.array
+    ) -> (mx.array, mx.array):
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        query = mx.transpose(mx.reshape(query, (1, -1, 24, 128)), (0, 2, 1, 3))
+        key = mx.transpose(mx.reshape(key, (1, -1, 24, 128)), (0, 2, 1, 3))
+        value = mx.transpose(mx.reshape(value, (1, -1, 24, 128)), (0, 2, 1, 3))
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        query, key = SingleBlockAttention.apply_rope(query, key, image_rotary_emb)
+
+        hidden_states = SingleBlockAttention.attention(query, key, value)
+        hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3))
+        hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension))
+
+        return hidden_states
+
+    @staticmethod
+    def attention(query, key, value):
+        scale = 1 / mx.sqrt(query.shape[-1])
+        scores = (query * scale) @ key.transpose(0, 1, 3, 2)
+        attn = mx.softmax(scores, axis=-1)
+        hidden_states = (attn @ value)
+        return hidden_states
+
+    @staticmethod
+    def apply_rope(xq: mx.array, xk: mx.array, freqs_cis: mx.array):
+        xq_ = xq.astype(mx.float32).reshape(*xq.shape[:-1], -1, 1, 2)
+        xk_ = xk.astype(mx.float32).reshape(*xk.shape[:-1], -1, 1, 2)
+        xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+        xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+        return xq_out.reshape(*xq.shape).astype(mx.float32), xk_out.reshape(*xk.shape).astype(mx.float32)

src/flux_1_schnell/models/transformer/single_transformer_block.py

@@ -0,0 +1,35 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.transformer.ada_layer_norm_zero_single import AdaLayerNormZeroSingle
+from flux_1_schnell.models.transformer.single_block_attention import SingleBlockAttention
+
+
+class SingleTransformerBlock(nn.Module):
+
+    def __init__(self, layer):
+        super().__init__()
+        self.layer = layer
+        self.norm = AdaLayerNormZeroSingle()
+        self.proj_mlp = nn.Linear(3072, 4*3072)
+        self.attn = SingleBlockAttention()
+        self.proj_out = nn.Linear(3072 + 4*3072, 3072)
+
+    def forward(
+            self,
+            hidden_states: mx.array,
+            text_embeddings: mx.array,
+            rotary_embeddings: mx.array
+    ) -> (mx.array, mx.array):
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm.forward(x=hidden_states, text_embeddings=text_embeddings)
+        mlp_hidden_states = nn.gelu_approx(self.proj_mlp(norm_hidden_states))
+        attn_output = self.attn.forward(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=rotary_embeddings,
+        )
+        hidden_states = mx.concatenate([attn_output, mlp_hidden_states], axis=2)
+        gate = mx.expand_dims(gate, axis=1)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states

src/flux_1_schnell/models/transformer/text_embedder.py

@@ -0,0 +1,16 @@
+from mlx import nn
+import mlx.core as mx
+
+
+class TextEmbedder(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.linear_1 = nn.Linear(768, 3072)
+        self.linear_2 = nn.Linear(3072, 3072)
+
+    def forward(self, caption: mx.array) -> mx.array:
+        hidden_states = self.linear_1(caption)
+        hidden_states = nn.silu(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states

src/flux_1_schnell/models/transformer/time_text_embed.py

@@ -0,0 +1,35 @@
+import math
+from mlx import nn
+import mlx.core as mx
+
+from flux_1_schnell.config.config import Config
+from flux_1_schnell.models.transformer.text_embedder import TextEmbedder
+from flux_1_schnell.models.transformer.timestep_embedder import TimestepEmbedder
+
+
+class TimeTextEmbed(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.text_embedder = TextEmbedder()
+        self.timestep_embedder = TimestepEmbedder()
+
+    def forward(self, time_step: mx.array, pooled_projection: mx.array) -> mx.array:
+        time_steps_proj = TimeTextEmbed._time_proj(time_step)
+        time_steps_emb = self.timestep_embedder.forward(time_steps_proj)
+        pooled_projections = self.text_embedder.forward(pooled_projection)
+        conditioning = time_steps_emb + pooled_projections
+        return conditioning.astype(Config.precision)
+
+    @staticmethod
+    def _time_proj(time_steps: mx.array) -> mx.array:
+        max_period = 10000
+        half_dim = 128
+        exponent = -math.log(max_period) * mx.arange(start=0, stop=half_dim, step=None, dtype=mx.float32)
+        exponent = exponent / half_dim
+        emb = mx.exp(exponent)
+        emb = time_steps[:, None].astype(mx.float32) * emb[None, :]
+        emb = mx.concatenate([mx.sin(emb), mx.cos(emb)], axis=-1)
+        emb = mx.concatenate([emb[:, half_dim:], emb[:, :half_dim]], axis=-1)
+        return emb
+

src/flux_1_schnell/models/transformer/timestep_embedder.py

@@ -0,0 +1,16 @@
+from mlx import nn
+import mlx.core as mx
+
+
+class TimestepEmbedder(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.linear_1 = nn.Linear(256, 3072)
+        self.linear_2 = nn.Linear(3072, 3072)
+
+    def forward(self, sample: mx.array) -> mx.array:
+        sample = self.linear_1(sample)
+        sample = nn.silu(sample)
+        sample = self.linear_2(sample)
+        return sample

src/flux_1_schnell/models/transformer/transformer.py

@@ -0,0 +1,82 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.config.config import Config
+from flux_1_schnell.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
+from flux_1_schnell.models.transformer.embed_nd import EmbedND
+from flux_1_schnell.models.transformer.joint_transformer_block import JointTransformerBlock
+from flux_1_schnell.models.transformer.single_transformer_block import SingleTransformerBlock
+from flux_1_schnell.models.transformer.time_text_embed import TimeTextEmbed
+
+
+class Transformer(nn.Module):
+
+    def __init__(self, weights: dict):
+        super().__init__()
+        self.pos_embed = EmbedND()
+        self.x_embedder = nn.Linear(64, 3072)
+        self.time_text_embed = TimeTextEmbed()
+        self.context_embedder = nn.Linear(4096, 3072)
+        self.transformer_blocks = [JointTransformerBlock(i) for i in range(19)]
+        self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)]
+        self.norm_out = AdaLayerNormContinuous(3072, 3072)
+        self.proj_out = nn.Linear(3072, 64)
+
+        # Load the weights after all components are initialized
+        self.update(weights)
+
+    def predict(
+            self,
+            t: int,
+            prompt_embeds: mx.array,
+            pooled_prompt_embeds: mx.array,
+            hidden_states: mx.array,
+            config: Config
+    ) -> mx.array:
+        time_step = config.time_steps[t]
+        time_step = mx.broadcast_to(time_step, (1,))
+        hidden_states = self.x_embedder(hidden_states)
+        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()
+        ids = mx.concatenate((txt_ids, img_ids), axis=1)
+        image_rotary_emb = self.pos_embed.forward(ids)
+
+        for block in self.transformer_blocks:
+            encoder_hidden_states, hidden_states = block.forward(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                text_embeddings=text_embeddings,
+                rotary_embeddings=image_rotary_emb
+            )
+
+        hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+        for block in self.single_transformer_blocks:
+            hidden_states = block.forward(
+                hidden_states=hidden_states,
+                text_embeddings=text_embeddings,
+                rotary_embeddings=image_rotary_emb
+            )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+        hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
+        hidden_states = self.proj_out(hidden_states)
+        noise = hidden_states
+        mx.eval(noise)
+        print(t)
+        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, :])
+        latent_image_ids = mx.repeat(latent_image_ids[None, :], 1, axis=0)
+        latent_image_ids = mx.reshape(latent_image_ids, (1, 4096, 3))
+        return latent_image_ids
+
+    @staticmethod
+    def _prepare_text_ids() -> mx.array:
+        return mx.zeros((1, 256, 3))

src/flux_1_schnell/models/vae/common/attention.py

@@ -0,0 +1,36 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.config.config import Config
+
+
+class Attention(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.group_norm = nn.GroupNorm(32, 512, pytorch_compatible=True)
+        self.to_q = nn.Linear(512, 512)
+        self.to_k = nn.Linear(512, 512)
+        self.to_v = nn.Linear(512, 512)
+        self.to_out = [nn.Linear(512, 512)]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+
+        B, H, W, C = input_array.shape
+
+        y = self.group_norm(input_array.astype(mx.float32)).astype(Config.precision)
+
+        queries = self.to_q(y).reshape(B, H * W, C)
+        keys = self.to_k(y).reshape(B, H * W, C)
+        values = self.to_v(y).reshape(B, H * W, C)
+
+        scale = 1 / mx.sqrt(queries.shape[-1])
+        scores = (queries * scale) @ keys.transpose(0, 2, 1)
+        attn = mx.softmax(scores, axis=-1)
+        y = (attn @ values).reshape(B, H, W, C)
+
+        y = self.to_out[0](y)
+        output_tensor = input_array + y
+
+        return mx.transpose(output_tensor, (0, 3, 1, 2))

src/flux_1_schnell/models/vae/common/resnet_block_2d.py

@@ -0,0 +1,69 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.config.config import Config
+
+
+class ResnetBlock2D(nn.Module):
+
+    def __init__(
+            self,
+            norm1: int,
+            conv1_in: int,
+            conv1_out: int,
+            norm2: int,
+            conv2_in: int,
+            conv2_out: int,
+            conv_shortcut_in: int | None = None,
+            conv_shortcut_out: int | None = None,
+            is_conv_shortcut: bool = False
+    ):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(
+            num_groups=32,
+            dims=norm1,
+            eps=1e-6,
+            affine=True,
+            pytorch_compatible=True
+        )
+        self.norm2 = nn.GroupNorm(
+            num_groups=32,
+            dims=norm2,
+            eps=1e-6,
+            affine=True,
+            pytorch_compatible=True
+        )
+        self.conv1 = nn.Conv2d(
+            in_channels=conv1_in,
+            out_channels=conv1_out,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1),
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels=conv2_in,
+            out_channels=conv2_out,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1),
+        )
+        self.is_conv_shortcut = is_conv_shortcut
+        self.conv_shortcut = None if not is_conv_shortcut else nn.Conv2d(
+            in_channels=conv_shortcut_in,
+            out_channels=conv_shortcut_out,
+            kernel_size=(1, 1),
+            stride=(1, 1),
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+        hidden_states = self.norm1(input_array.astype(mx.float32)).astype(Config.precision)
+        hidden_states = nn.silu(hidden_states)
+        hidden_states = self.conv1(hidden_states)
+        hidden_states = self.norm2(hidden_states.astype(mx.float32)).astype(Config.precision)
+        hidden_states = nn.silu(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+        if self.is_conv_shortcut:
+            input_array = self.conv_shortcut(input_array)
+        output_tensor = input_array + hidden_states
+        return mx.transpose(output_tensor, (0, 3, 1, 2))

src/flux_1_schnell/models/vae/common/unet_mid_block.py

@@ -0,0 +1,22 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.attention import Attention
+from flux_1_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
+
+
+class UnetMidBlock(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.attentions = [Attention()]
+        self.resnets = [
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512)
+        ]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = self.resnets[0].forward(input_array)
+        hidden_states = self.attentions[0].forward(hidden_states)
+        hidden_states = self.resnets[1].forward(hidden_states)
+        return hidden_states

src/flux_1_schnell/models/vae/decoder/conv_in.py

@@ -0,0 +1,19 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class ConvIn(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv2d = nn.Conv2d(
+            in_channels=16,
+            out_channels=512,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+        return mx.transpose(self.conv2d(input_array), (0, 3, 1, 2))

src/flux_1_schnell/models/vae/decoder/conv_norm_out.py

@@ -0,0 +1,21 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from flux_1_schnell.config.config import Config
+
+
+class ConvNormOut(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm = nn.GroupNorm(
+            num_groups=32,
+            dims=128,
+            eps=1e-6,
+            affine=True,
+            pytorch_compatible=True
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+        hidden_states = self.norm(input_array.astype(mx.float32)).astype(Config.precision)
+        return mx.transpose(hidden_states, (0, 3, 1, 2))

src/flux_1_schnell/models/vae/decoder/conv_out.py

@@ -0,0 +1,18 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class ConvOut(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv2d = nn.Conv2d(
+            in_channels=128,
+            out_channels=3,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1),
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+        return mx.transpose(self.conv2d(input_array), (0, 3, 1, 2))

src/flux_1_schnell/models/vae/decoder/decoder.py

@@ -0,0 +1,36 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.decoder.conv_in import ConvIn
+from flux_1_schnell.models.vae.decoder.conv_norm_out import ConvNormOut
+from flux_1_schnell.models.vae.decoder.conv_out import ConvOut
+from flux_1_schnell.models.vae.common.unet_mid_block import UnetMidBlock
+from flux_1_schnell.models.vae.decoder.up_block_1_or_2 import UpBlock1Or2
+from flux_1_schnell.models.vae.decoder.up_block_3 import UpBlock3
+from flux_1_schnell.models.vae.decoder.up_block_4 import UpBlock4
+
+
+class Decoder(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv_in = ConvIn()
+        self.mid_block = UnetMidBlock()
+        self.up_blocks = [
+            UpBlock1Or2(),
+            UpBlock1Or2(),
+            UpBlock3(),
+            UpBlock4(),
+        ]
+        self.conv_norm_out = ConvNormOut()
+        self.conv_out = ConvOut()
+
+    def decode(self, latents: mx.array) -> mx.array:
+        latents = self.conv_in.forward(latents)
+        latents = self.mid_block.forward(latents)
+        for up_block in self.up_blocks:
+            latents = up_block.forward(latents)
+        latents = self.conv_norm_out.forward(latents)
+        latents = nn.silu(latents)
+        latents = self.conv_out.forward(latents)
+        return latents

src/flux_1_schnell/models/vae/decoder/up_block_1_or_2.py

@@ -0,0 +1,27 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1_schnell.models.vae.decoder.up_sampler import UpSampler
+
+
+class UpBlock1Or2(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.resnets = [
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512)
+        ]
+        self.upsamplers = [UpSampler(conv_in=512, conv_out=512)]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = self.resnets[0].forward(input_array)
+        hidden_states = self.resnets[1].forward(hidden_states)
+        hidden_states = self.resnets[2].forward(hidden_states)
+
+        if self.upsamplers is not None:
+            hidden_states = self.upsamplers[0].forward(hidden_states)
+
+        return hidden_states

src/flux_1_schnell/models/vae/decoder/up_block_3.py

@@ -0,0 +1,27 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1_schnell.models.vae.decoder.up_sampler import UpSampler
+
+
+class UpBlock3(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.resnets = [
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256, is_conv_shortcut=True, conv_shortcut_in=512, conv_shortcut_out=256),
+            ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256),
+            ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256),
+        ]
+        self.upsamplers = [UpSampler(conv_in=256, conv_out=256)]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = self.resnets[0].forward(input_array)
+        hidden_states = self.resnets[1].forward(hidden_states)
+        hidden_states = self.resnets[2].forward(hidden_states)
+
+        if self.upsamplers is not None:
+            hidden_states = self.upsamplers[0].forward(hidden_states)
+
+        return hidden_states

src/flux_1_schnell/models/vae/decoder/up_block_4.py

@@ -0,0 +1,21 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
+
+
+class UpBlock4(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.resnets = [
+            ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128, is_conv_shortcut=True, conv_shortcut_in=256, conv_shortcut_out=128),
+            ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128),
+            ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128),
+        ]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = self.resnets[0].forward(input_array)
+        hidden_states = self.resnets[1].forward(hidden_states)
+        hidden_states = self.resnets[2].forward(hidden_states)
+        return hidden_states

src/flux_1_schnell/models/vae/decoder/up_sampler.py

@@ -0,0 +1,29 @@
+import mlx.core as mx
+
+from mlx import nn
+
+
+class UpSampler(nn.Module):
+
+    def __init__(self, conv_in: int, conv_out: int):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels=conv_in,
+            out_channels=conv_out,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1),
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+        hidden_states = UpSampler.up_sample_nearest(input_array)
+        hidden_state = self.conv(hidden_states)
+        return mx.transpose(hidden_state, (0, 3, 1, 2))
+
+    @staticmethod
+    def up_sample_nearest(x: mx.array, scale: int = 2):
+        B, H, W, C = x.shape
+        x = mx.broadcast_to(x[:, :, None, :, None, :], (B, H, scale, W, scale, C))
+        x = x.reshape(B, H * scale, W * scale, C)
+        return x

src/flux_1_schnell/models/vae/encoder/conv_in.py

@@ -0,0 +1,19 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class ConvIn(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv2d = nn.Conv2d(
+            in_channels=3,
+            out_channels=128,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+        return mx.transpose(self.conv2d(input_array), (0, 3, 1, 2))

src/flux_1_schnell/models/vae/encoder/conv_norm_out.py

@@ -0,0 +1,19 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class ConvNormOut(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm = nn.GroupNorm(
+            num_groups=32,
+            dims=512,
+            eps=1e-6,
+            affine=True,
+            pytorch_compatible=True
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+        hidden_states = self.norm(input_array.astype(mx.float32)).astype(mx.float32)
+        return mx.transpose(hidden_states, (0, 3, 1, 2))

src/flux_1_schnell/models/vae/encoder/conv_out.py

@@ -0,0 +1,18 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class ConvOut(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv2d = nn.Conv2d(
+            in_channels=512,
+            out_channels=32,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1),
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        input_array = mx.transpose(input_array, (0, 2, 3, 1))
+        return mx.transpose(self.conv2d(input_array), (0, 3, 1, 2))

src/flux_1_schnell/models/vae/encoder/down_block_1.py

@@ -0,0 +1,25 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1_schnell.models.vae.encoder.down_sampler import DownSampler
+
+
+class DownBlock1(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.resnets = [
+            ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128),
+            ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=128, norm2=128, conv2_in=128, conv2_out=128),
+        ]
+        self.downsamplers = [DownSampler(conv_in=128, conv_out=128)]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = self.resnets[0].forward(input_array)
+        hidden_states = self.resnets[1].forward(hidden_states)
+
+        if self.downsamplers is not None:
+            hidden_states = self.downsamplers[0].forward(hidden_states)
+
+        return hidden_states

src/flux_1_schnell/models/vae/encoder/down_block_2.py

@@ -0,0 +1,25 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1_schnell.models.vae.encoder.down_sampler import DownSampler
+
+
+class DownBlock2(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.resnets = [
+            ResnetBlock2D(norm1=128, conv1_in=128, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256, is_conv_shortcut=True, conv_shortcut_in=128, conv_shortcut_out=256),
+            ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=256, norm2=256, conv2_in=256, conv2_out=256),
+        ]
+        self.downsamplers = [DownSampler(conv_in=256, conv_out=256)]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = self.resnets[0].forward(input_array)
+        hidden_states = self.resnets[1].forward(hidden_states)
+
+        if self.downsamplers is not None:
+            hidden_states = self.downsamplers[0].forward(hidden_states)
+
+        return hidden_states

src/flux_1_schnell/models/vae/encoder/down_block_3.py

@@ -0,0 +1,25 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1_schnell.models.vae.encoder.down_sampler import DownSampler
+
+
+class DownBlock3(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.resnets = [
+            ResnetBlock2D(norm1=256, conv1_in=256, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512, is_conv_shortcut=True, conv_shortcut_in=256, conv_shortcut_out=512),
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
+        ]
+        self.downsamplers = [DownSampler(conv_in=512, conv_out=512)]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = self.resnets[0].forward(input_array)
+        hidden_states = self.resnets[1].forward(hidden_states)
+
+        if self.downsamplers is not None:
+            hidden_states = self.downsamplers[0].forward(hidden_states)
+
+        return hidden_states

src/flux_1_schnell/models/vae/encoder/down_block_4.py

@@ -0,0 +1,19 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
+
+
+class DownBlock4(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.resnets = [
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
+            ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
+        ]
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = self.resnets[0].forward(input_array)
+        hidden_states = self.resnets[1].forward(hidden_states)
+        return hidden_states

src/flux_1_schnell/models/vae/encoder/down_sampler.py

@@ -0,0 +1,21 @@
+import mlx.core as mx
+
+from mlx import nn
+
+
+class DownSampler(nn.Module):
+
+    def __init__(self, conv_in: int, conv_out: int):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels=conv_in,
+            out_channels=conv_out,
+            kernel_size=(3, 3),
+            stride=(2, 2),
+        )
+
+    def forward(self, input_array: mx.array) -> mx.array:
+        hidden_states = mx.pad(input_array, ((0, 0), (0, 0), (0, 1), (0, 1)))
+        hidden_states = mx.transpose(hidden_states, (0, 2, 3, 1))
+        hidden_state = self.conv(hidden_states)
+        return mx.transpose(hidden_state, (0, 3, 1, 2))

src/flux_1_schnell/models/vae/encoder/encoder.py

@@ -0,0 +1,38 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.common.unet_mid_block import UnetMidBlock
+
+from flux_1_schnell.models.vae.encoder.conv_in import ConvIn
+from flux_1_schnell.models.vae.encoder.conv_norm_out import ConvNormOut
+from flux_1_schnell.models.vae.encoder.conv_out import ConvOut
+from flux_1_schnell.models.vae.encoder.down_block_1 import DownBlock1
+from flux_1_schnell.models.vae.encoder.down_block_2 import DownBlock2
+from flux_1_schnell.models.vae.encoder.down_block_3 import DownBlock3
+from flux_1_schnell.models.vae.encoder.down_block_4 import DownBlock4
+
+
+class Encoder(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv_in = ConvIn()
+        self.mid_block = UnetMidBlock()
+        self.down_blocks = [
+            DownBlock1(),
+            DownBlock2(),
+            DownBlock3(),
+            DownBlock4(),
+        ]
+        self.conv_norm_out = ConvNormOut()
+        self.conv_out = ConvOut()
+
+    def encode(self, latents: mx.array) -> mx.array:
+        latents = self.conv_in.forward(latents)
+        for down_block in self.down_blocks:
+            latents = down_block.forward(latents)
+        latents = self.mid_block.forward(latents)
+        latents = self.conv_norm_out.forward(latents)
+        latents = nn.silu(latents)
+        latents = self.conv_out.forward(latents)
+        return latents

src/flux_1_schnell/models/vae/vae.py

@@ -0,0 +1,27 @@
+import mlx.core as mx
+from mlx import nn
+
+from flux_1_schnell.models.vae.decoder.decoder import Decoder
+from flux_1_schnell.models.vae.encoder.encoder import Encoder
+
+
+class VAE(nn.Module):
+    scaling_factor: int = 0.3611
+    shift_factor: int = 0.1159
+
+    def __init__(self, weights: dict):
+        super().__init__()
+        self.decoder = Decoder()
+        self.encoder = Encoder()
+
+        # Load the weights after all components are initialized
+        self.update(weights)
+
+    def decode(self, latents: mx.array) -> mx.array:
+        scaled_latents = (latents / self.scaling_factor) + self.shift_factor
+        return self.decoder.decode(scaled_latents)
+
+    def encode(self, latents: mx.array) -> mx.array:
+        latents = self.encoder.encode(latents)
+        mean, _ = mx.split(latents, 2, axis=1)
+        return (mean - self.shift_factor) * self.scaling_factor

src/flux_1_schnell/post_processing/image_util.py

@@ -0,0 +1,55 @@
+import PIL
+import mlx.core as mx
+import numpy as np
+from PIL import Image
+
+
+class ImageUtil:
+
+    @staticmethod
+    def to_image(decoded_latents: mx.array) -> PIL.Image.Image:
+        normalized = ImageUtil._denormalize(decoded_latents)
+        normalized_numpy = ImageUtil._to_numpy(normalized)
+        image = ImageUtil._numpy_to_pil(normalized_numpy)
+        return image
+
+    @staticmethod
+    def _denormalize(images: mx.array) -> mx.array:
+        return mx.clip((images / 2 + 0.5), 0, 1)
+
+    @staticmethod
+    def _normalize(images: mx.array) -> mx.array:
+        return 2.0 * images - 1.0
+
+    @staticmethod
+    def _to_numpy(images: mx.array) -> np.ndarray:
+        images = mx.transpose(images, (0, 2, 3, 1))
+        images = mx.array.astype(images, mx.float32)
+        images = np.array(images)
+        return images
+
+    @staticmethod
+    def _numpy_to_pil(images: np.ndarray) -> PIL.Image.Image:
+        images = (images * 255).round().astype("uint8")
+        pil_images = [Image.fromarray(image) for image in images]
+        return pil_images[0]
+
+    @staticmethod
+    def _pil_to_numpy(image: PIL.Image.Image) -> np.ndarray:
+        image = np.array(image).astype(np.float32) / 255.0
+        images = np.stack([image], axis=0)
+        return images
+
+    @staticmethod
+    def to_array(image: PIL.Image.Image) -> mx.array:
+        image = ImageUtil.resize(image)
+        image = ImageUtil._pil_to_numpy(image)
+        array = mx.array(image)
+        array = mx.transpose(array, (0, 3, 1, 2))
+        array = ImageUtil._normalize(array)
+        return array
+
+    @staticmethod
+    def resize(image):
+        image = image.resize((1024, 1024), resample=PIL.Image.LANCZOS)
+        return image

src/flux_1_schnell/scheduler/scheduler.py

@@ -0,0 +1,20 @@
+import mlx.core as mx
+
+from flux_1_schnell.config.config import Config
+
+
+class FlowMatchEulerDiscreteNoiseScheduler:
+
+    @staticmethod
+    def denoise(
+            t: int,
+            noise: mx.array,
+            latent: mx.array,
+            config: Config,
+    ) -> mx.array:
+        sigma = config.sigmas[t]
+        denoised = latent - noise * sigma
+        derivative = (latent - denoised) / sigma
+        dt = config.sigmas[t + 1] - sigma
+        prev_sample = latent + derivative * dt
+        return prev_sample

src/flux_1_schnell/tokenizer/clip_tokenizer.py

@@ -0,0 +1,21 @@
+import mlx.core as mx
+from transformers import CLIPTokenizer
+
+
+class TokenizerCLIP:
+    MAX_TOKEN_LENGTH = 77
+
+    def __init__(self, tokenizer: CLIPTokenizer):
+        self.tokenizer = tokenizer
+
+    def tokenize(self, prompt: str) -> mx.array:
+        text_input_ids = self.tokenizer(
+            [prompt],
+            padding="max_length",
+            max_length=TokenizerCLIP.MAX_TOKEN_LENGTH,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        ).input_ids
+        return mx.array(text_input_ids.cpu().numpy())

src/flux_1_schnell/tokenizer/t5_tokenizer.py

@@ -0,0 +1,21 @@
+import mlx.core as mx
+from transformers import T5Tokenizer
+
+
+class TokenizerT5:
+    MAX_TOKEN_LENGTH = 256
+
+    def __init__(self, tokenizer: T5Tokenizer):
+        self.tokenizer = tokenizer
+
+    def tokenize(self, prompt: str) -> mx.array:
+        text_input_ids = self.tokenizer(
+            [prompt],
+            padding="max_length",
+            max_length=TokenizerT5.MAX_TOKEN_LENGTH,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        ).input_ids
+        return mx.array(text_input_ids.cpu().numpy())

src/flux_1_schnell/tokenizer/tokenizer_handler.py

@@ -0,0 +1,23 @@
+import transformers
+
+from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
+from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
+
+
+class TokenizerHandler:
+
+    def __init__(self, root_path: str):
+        self.clip = transformers.CLIPTokenizer.from_pretrained(
+            pretrained_model_name_or_path=root_path + "/tokenizer",
+            local_files_only=True,
+            max_length=TokenizerCLIP.MAX_TOKEN_LENGTH
+        )
+        self.t5 = transformers.T5Tokenizer.from_pretrained(
+            pretrained_model_name_or_path=root_path + "/tokenizer_2",
+            local_files_only=True,
+            max_length=TokenizerT5.MAX_TOKEN_LENGTH
+        )
+
+    @staticmethod
+    def load_from_disk_via_huggingface_transformers(root_path: str) -> "TokenizerHandler":
+        return TokenizerHandler(root_path)

src/flux_1_schnell/weights/weight_handler.py

@@ -0,0 +1,110 @@
+import mlx.core as mx
+from mlx.utils import tree_unflatten
+
+from flux_1_schnell.config.config import Config
+
+
+class WeightHandler:
+
+    def __init__(self, root_path: str):
+        self.clip_encoder = WeightHandler._clip_encoder(
+            weights=WeightHandler._load(root_path + "/text_encoder/model.safetensors")
+        )
+        self.t5_encoder = WeightHandler._t5_encoder(
+            weights_1=WeightHandler._load(root_path + "/text_encoder_2/model-00001-of-00002.safetensors"),
+            weights_2=WeightHandler._load(root_path + "/text_encoder_2/model-00002-of-00002.safetensors"),
+        )
+        self.vae = WeightHandler._vae(
+            weights=WeightHandler._load(root_path + "/vae/diffusion_pytorch_model.safetensors")
+        )
+        self.transformer = WeightHandler._transformer(
+            weights_1=WeightHandler._load(root_path + "transformer/diffusion_pytorch_model-00001-of-00003.safetensors"),
+            weights_2=WeightHandler._load(root_path + "transformer/diffusion_pytorch_model-00002-of-00003.safetensors"),
+            weights_3=WeightHandler._load(root_path + "transformer/diffusion_pytorch_model-00003-of-00003.safetensors"),
+        )
+
+    @staticmethod
+    def load_from_disk(root_path: str) -> "WeightHandler":
+        return WeightHandler(root_path)
+
+    @staticmethod
+    def _load(path: str) -> list[dict]:
+        return list(mx.load(path).items())
+
+    @staticmethod
+    def _clip_encoder(weights: list[dict]) -> dict:
+        weights = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights])
+        unflatten = tree_unflatten(weights)
+        return unflatten
+
+    @staticmethod
+    def _t5_encoder(
+            weights_1: list[dict],
+            weights_2: list[dict],
+    ) -> dict:
+        weights_1 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_1])
+        weights_2 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_2])
+        unflatten = tree_unflatten(weights_1 + weights_2)
+        unflatten["final_layer_norm"] = unflatten["encoder"]["final_layer_norm"]
+        for block in unflatten["encoder"]["block"]:
+            attention = block["layer"][0]
+            ff = block["layer"][1]
+            block.pop("layer")
+            block["attention"] = attention
+            block["ff"] = ff
+
+        unflatten["t5_blocks"] = unflatten["encoder"]["block"]
+
+        # Only the first layer has the weights for "relative_attention_bias", we duplicate them here to keep code simple
+        relative_attention_bias = unflatten["t5_blocks"][0]["attention"]["SelfAttention"]["relative_attention_bias"]
+        for block in unflatten["t5_blocks"][1:]:
+            block["attention"]["SelfAttention"]["relative_attention_bias"] = relative_attention_bias
+
+        unflatten.pop("encoder")
+        return unflatten
+
+    @staticmethod
+    def _transformer(
+            weights_1: list[dict],
+            weights_2: list[dict],
+            weights_3: list[dict],
+    ) -> dict:
+        weights_1 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_1])
+        weights_2 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_2])
+        weights_3 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_3])
+        unflatten = tree_unflatten(weights_1 + weights_2 + weights_3)
+
+        for block in unflatten["transformer_blocks"]:
+            block["ff"] = {
+                "linear1": block["ff"]["net"][0]["proj"],
+                "linear2": block["ff"]["net"][2]
+            }
+            if block.get("ff_context") is not None:
+                block["ff_context"] = {
+                    "linear1": block["ff_context"]["net"][0]["proj"],
+                    "linear2": block["ff_context"]["net"][2]
+                }
+        return unflatten
+
+    @staticmethod
+    def _vae(weights: list[dict]) -> dict:
+        weights = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights])
+        unflatten = tree_unflatten(weights)
+        unflatten['decoder']['conv_in'] = {'conv2d': unflatten['decoder']['conv_in']}
+        unflatten['decoder']['conv_out'] = {'conv2d': unflatten['decoder']['conv_out']}
+        unflatten['decoder']['conv_norm_out'] = {'norm': unflatten['decoder']['conv_norm_out']}
+        unflatten['encoder']['conv_in'] = {'conv2d': unflatten['encoder']['conv_in']}
+        unflatten['encoder']['conv_out'] = {'conv2d': unflatten['encoder']['conv_out']}
+        unflatten['encoder']['conv_norm_out'] = {'norm': unflatten['encoder']['conv_norm_out']}
+        return unflatten
+
+    @staticmethod
+    def _flatten(params):
+        return [(k, v) for p in params for (k, v) in p]
+
+    @staticmethod
+    def _reshape_weights(key, value):
+        if len(value.shape) == 4:
+            value = value.transpose(0, 2, 3, 1)
+        value = value.reshape(-1).reshape(value.shape).astype(Config.precision)
+        return [(key, value)]