Merge pull request #32 from filipstrand/PyPI-release

PyPI Release

README.md

@@ -1,5 +1,5 @@
 
-![image](src/flux_1/assets/logo.png)
+![image](src/mflux/assets/logo.png)
 *A MLX port of FLUX based on the Huggingface Diffusers implementation.*
 
 ### About
@@ -23,30 +23,39 @@ like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) f
 - [x] FLUX.1-Dev
 
 ### Installation
+For users, the easiest way to install MFLUX is via pip:
+   ```
+   pip install mflux
+   ```
+<details>
+<summary>For contributors (click to expand)</summary>
+
 1. Clone the repo:
-    ```
-    git clone git@github.com:filipstrand/mflux.git
-    ```
+ ```
+ git clone git@github.com:filipstrand/mflux.git
+ ```
 2. Navigate to the project and set up a virtual environment:
-   ```
-   cd mflux && python3 -m venv .venv && source .venv/bin/activate
-   ``` 
+ ```
+ cd mflux && python3 -m venv .venv && source .venv/bin/activate
+ ```
 3. Install the required dependencies:
-    ```
-   pip install -r requirements.txt
-   ```
+ ```
+ pip install -r requirements.txt
+ ```
+</details>
+
 ### Generating an image
 
-Run the provided [main.py](main.py) by specifying a prompt and some optional arguments like so using the `schnell` model:
+Run the command `mflux-generate` by specifying a prompt and the model and some optional arguments. For example, here we use a quantized version of the `schnell` model for 2 steps:
 
 ```
-python main.py --model schnell --prompt "Luxury food photograph" --steps 2 --seed 2 -q 8
+mflux-generate --model schnell --prompt "Luxury food photograph" --steps 2 --seed 2 -q 8
 ```
 
-or use the slower, but more powerful `dev` model and run it with more time steps:
+This example uses the more powerful `dev` model with 25 time steps:
 
 ```
-python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 -q 8
+mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 -q 8
 ```
 
 ⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#quantization) section for running compressed versions of the model.* ⚠️
@@ -89,18 +98,17 @@ python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
  
 - **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `exiftool image.png`)
 
-Or, with the correct python environment active, make a new separate script like the following:
+Or, with the correct python environment active, create and run a separate script like the following:
 
 ```python
-import sys
-
-sys.path.append("/path/to/mflux/src")
-
-from flux_1.config.config import Config
-from flux_1.flux import Flux1
+from mflux.flux.flux import Flux1
+from mflux.config.config import Config
 
 # Load the model
-flux = Flux1.from_alias(alias="schnell")  # "schnell" or "dev"
+flux = Flux1.from_alias(
+   alias="schnell",  # "schnell" or "dev"
+   quantize=8,       # 4 or 8
+)  
 
 # Generate an image
 image = flux.generate_image(
@@ -116,14 +124,14 @@ image = flux.generate_image(
 image.save(path="image.png")
 ```
 
-For more options on how to configure MFLUX, please see [main.py](main.py).
+For more options on how to configure MFLUX, please see [generate.py](src/mflux/generate.py).
 
 ### Image generation speed (updated)
 
 These numbers are based on the non-quantized `schnell` model, with the configuration provided in the code snippet below. 
 To time your machine, run the following:
 ```
-time python main.py \
+time mflux-generate \
 --prompt "Luxury food photograph" \
 --model schnell \
 --steps 2 \
@@ -157,27 +165,27 @@ However, if we were to import a fixed instance of this latent array saved from t
 The images below illustrate this equivalence. 
 In all 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/config/config.py) class. 
+The precision for MFLUX can be set in the [Config](src/mflux/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/assets/comparison1.jpg)
+![image](src/mflux/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/assets/comparison2.jpg)
+![image](src/mflux/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/assets/comparison3.jpg)
+![image](src/mflux/assets/comparison3.jpg)
 
 ---
 
@@ -185,14 +193,14 @@ photorealistic, lotr, A tiny red dragon curled up asleep inside a nest, (Soft Fo
 ```
 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/assets/comparison4.jpg)
+![image](src/mflux/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/assets/comparison5.jpg)
+![image](src/mflux/assets/comparison5.jpg)
 
 ---
 
@@ -202,7 +210,7 @@ MFLUX supports running FLUX in 4-bit or 8-bit quantized mode. Running a quantize
 generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#size-comparisons-for-quantized-models). 
 
 ```
-python main.py \
+mflux-generate \
     --model schnell \
     --steps 2 \
     --seed 2 \
@@ -211,7 +219,7 @@ python main.py \
     --width 1024 \
     --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier" 
 ```
-![image](src/flux_1/assets/comparison6.jpg)
+![image](src/mflux/assets/comparison6.jpg)
 
 *In this example, weights are quantized at **runtime** - this is convenient if you don't want to [save a quantized copy of the weights to disk](#saving-a-quantized-version-to-disk), but still want to benefit from the potential speedup and RAM reduction quantization might bring.*
 
@@ -232,10 +240,10 @@ The reason weights sizes are not fully cut in half is because a small number of
 
 #### Saving a quantized version to disk
 
-To save a local copy of the quantized weights, run the `save.py` script like so:
+To save a local copy of the quantized weights, run the `mflux-save` command like so:
 
 ```
-python save.py \
+mflux-save \
     --path "/Users/filipstrand/Desktop/schnell_8bit" \
     --model schnell \
     --quantize 8
@@ -248,7 +256,7 @@ python save.py \
 To generate a new image from the quantized model, simply provide a `--path` to where it was saved: 
 
 ```
-python main.py \
+mflux-generate \
     --path "/Users/filipstrand/Desktop/schnell_8bit" \
     --model schnell \
     --steps 2 \
@@ -268,7 +276,7 @@ MFLUX also supports running a non-quantized model directly from a custom locatio
 In the example below, the model is placed in `/Users/filipstrand/Desktop/schnell`:
 
 ```
-python main.py \
+mflux-generate \
     --path "/Users/filipstrand/Desktop/schnell" \
     --model schnell \
     --steps 2 \
@@ -317,18 +325,18 @@ MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/tr
 The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) LoRA from [@TheLastBen](https://github.com/TheLastBen): 
 
 ```
-python main.py --prompt "sandor clegane" --model dev --steps 20 --seed 43 -q 8 --lora-paths "sandor_clegane_single_layer.safetensors"
+mflux-generate --prompt "sandor clegane" --model dev --steps 20 --seed 43 -q 8 --lora-paths "sandor_clegane_single_layer.safetensors"
 ```
 
-![image](src/flux_1/assets/lora1.jpg)
+![image](src/mflux/assets/lora1.jpg)
 ---
 
 The following example is [Flux_1_Dev_LoRA_Paper-Cutout-Style](https://huggingface.co/Norod78/Flux_1_Dev_LoRA_Paper-Cutout-Style) LoRA from [@Norod78](https://huggingface.co/Norod78):
 
 ```
-python main.py --prompt "pikachu, Paper Cutout Style" --model schnell --steps 4 --seed 43 -q 8 --lora-paths "Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors"
+mflux-generate --prompt "pikachu, Paper Cutout Style" --model schnell --steps 4 --seed 43 -q 8 --lora-paths "Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors"
 ```
-![image](src/flux_1/assets/lora2.jpg)
+![image](src/mflux/assets/lora2.jpg)
 
 *Note that LoRA trained weights are typically trained with a **trigger word or phrase**. For example, in the latter case, the sentence should include the phrase **"Paper Cutout Style"**.*
 
@@ -339,7 +347,7 @@ python main.py --prompt "pikachu, Paper Cutout Style" --model schnell --steps 4
 Multiple LoRAs can be sent in to combine the effects of the individual adapters. The following example combines both of the above LoRAs:
 
 ```
-python main.py \
+mflux-generate \
    --prompt "sandor clegane in a forest, Paper Cutout Style" \
    --model dev \
    --steps 20 \
@@ -348,7 +356,7 @@ python main.py \
    --lora-scales 1.0 1.0 \
    -q 8
 ```
-![image](src/flux_1/assets/lora3.jpg)
+![image](src/mflux/assets/lora3.jpg)
 
 Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all. 
 The example above also show the usage of `--lora-scales` flag. 

pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mflux"
+version = "0.2.0"
+description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
+readme = "README.md"
+authors = [
+    { name = "Filip Strand", email = "strand.filip@gmail.com" }
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: MacOS",
+]
+dependencies = [
+    "mlx>=0.16.0",
+    "numpy>=2.0.0",
+    "pillow>=10.4.0",
+    "transformers>=4.44.0",
+    "sentencepiece>=0.2.0",
+    "torch>=2.3.1",
+    "tqdm>=4.66.5",
+    "huggingface-hub>=0.24.5",
+    "safetensors>=0.4.4",
+    "piexif>=1.1.3",
+]
+
+[project.urls]
+homepage = "https://github.com/filipstrand/mflux"
+
+[project.scripts]
+mflux-generate = "mflux.generate:main"
+mflux-save = "mflux.save:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["mflux*"]

src/mflux/config/runtime_config.py

@@ -1,8 +1,8 @@
 import mlx.core as mx
 import numpy as np
 
-from flux_1.config.config import Config
-from flux_1.config.model_config import ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 
 
 class RuntimeConfig:

src/mflux/flux/flux.py

@@ -5,19 +5,19 @@ from mlx import nn
 from mlx.utils import tree_flatten
 from tqdm import tqdm
 
-from flux_1.config.config import Config
-from flux_1.config.model_config import ModelConfig
-from flux_1.config.runtime_config import RuntimeConfig
-from flux_1.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
-from flux_1.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
-from flux_1.models.transformer.transformer import Transformer
-from flux_1.models.vae.vae import VAE
-from flux_1.post_processing.image import Image
-from flux_1.post_processing.image_util import ImageUtil
-from flux_1.tokenizer.clip_tokenizer import TokenizerCLIP
-from flux_1.tokenizer.t5_tokenizer import TokenizerT5
-from flux_1.tokenizer.tokenizer_handler import TokenizerHandler
-from flux_1.weights.weight_handler import WeightHandler
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
+from mflux.models.transformer.transformer import Transformer
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.image import Image
+from mflux.post_processing.image_util import ImageUtil
+from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.tokenizer.t5_tokenizer import TokenizerT5
+from mflux.tokenizer.tokenizer_handler import TokenizerHandler
+from mflux.weights.weight_handler import WeightHandler
 
 
 class Flux1:
@@ -126,8 +126,11 @@ class Flux1:
         return latents
 
     @staticmethod
-    def from_alias(alias: str) -> "Flux1":
-        return Flux1(ModelConfig.from_alias(alias))
+    def from_alias(alias: str, quantize: int | None = None) -> "Flux1":
+        return Flux1(
+            model_config=ModelConfig.from_alias(alias),
+            quantize=quantize,
+        )
 
     def _set_model_weights(self, weights):
         self.vae.update(weights.vae)

src/mflux/generate.py

@@ -3,11 +3,11 @@ import os
 import sys
 import time
 
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from flux_1.config.model_config import ModelConfig
-from flux_1.config.config import Config
-from flux_1.flux import Flux1
+from mflux.config.model_config import ModelConfig
+from mflux.config.config import Config
+from mflux.flux.flux import Flux1
 
 
 def main():

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

@@ -1,7 +1,7 @@
 from mlx import nn
 import mlx.core as mx
 
-from flux_1.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
 
 
 class CLIPEmbeddings(nn.Module):

src/mflux/models/text_encoder/clip_encoder/clip_encoder.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.text_encoder.clip_encoder.clip_text_model import CLIPTextModel
+from mflux.models.text_encoder.clip_encoder.clip_text_model import CLIPTextModel
 
 
 class CLIPEncoder(nn.Module):

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

@@ -1,8 +1,8 @@
 from mlx import nn
 import mlx.core as mx
 
-from flux_1.models.text_encoder.clip_encoder.clip_mlp import CLIPMLP
-from flux_1.models.text_encoder.clip_encoder.clip_sdpa_attention import CLIPSdpaAttention
+from mflux.models.text_encoder.clip_encoder.clip_mlp import CLIPMLP
+from mflux.models.text_encoder.clip_encoder.clip_sdpa_attention import CLIPSdpaAttention
 
 
 class CLIPEncoderLayer(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.text_encoder.clip_encoder.clip_embeddings import CLIPEmbeddings
-from flux_1.models.text_encoder.clip_encoder.encoder_clip import EncoderCLIP
+from mflux.models.text_encoder.clip_encoder.clip_embeddings import CLIPEmbeddings
+from mflux.models.text_encoder.clip_encoder.encoder_clip import EncoderCLIP
 
 
 class CLIPTextModel(nn.Module):

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

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.text_encoder.clip_encoder.clip_encoder_layer import CLIPEncoderLayer
+from mflux.models.text_encoder.clip_encoder.clip_encoder_layer import CLIPEncoderLayer
 
 
 class EncoderCLIP(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
-from flux_1.models.text_encoder.t5_encoder.t5_self_attention import T5SelfAttention
+from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
+from mflux.models.text_encoder.t5_encoder.t5_self_attention import T5SelfAttention
 
 
 class T5Attention(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.text_encoder.t5_encoder.t5_attention import T5Attention
-from flux_1.models.text_encoder.t5_encoder.t5_feed_forward import T5FeedForward
+from mflux.models.text_encoder.t5_encoder.t5_attention import T5Attention
+from mflux.models.text_encoder.t5_encoder.t5_feed_forward import T5FeedForward
 
 
 class T5Block(nn.Module):

src/mflux/models/text_encoder/t5_encoder/t5_encoder.py

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.text_encoder.t5_encoder.t5_block import T5Block
-from flux_1.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
+from mflux.models.text_encoder.t5_encoder.t5_block import T5Block
+from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
 
 
 class T5Encoder(nn.Module):

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

@@ -3,8 +3,8 @@ import math
 from mlx import nn
 import mlx.core as mx
 
-from flux_1.models.text_encoder.t5_encoder.t5_dense_relu_dense import T5DenseReluDense
-from flux_1.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
+from mflux.models.text_encoder.t5_encoder.t5_dense_relu_dense import T5DenseReluDense
+from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
 
 
 class T5FeedForward(nn.Module):

src/mflux/models/transformer/ada_layer_norm_continous.py

@@ -1,7 +1,7 @@
 from mlx import nn
 import mlx.core as mx
 
-from flux_1.config.config import Config
+from mflux.config.config import Config
 
 
 class AdaLayerNormContinuous(nn.Module):

src/mflux/models/transformer/joint_transformer_block.py

@@ -1,9 +1,9 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.transformer.ada_layer_norm_zero import AdaLayerNormZero
-from flux_1.models.transformer.feed_forward import FeedForward
-from flux_1.models.transformer.joint_attention import JointAttention
+from mflux.models.transformer.ada_layer_norm_zero import AdaLayerNormZero
+from mflux.models.transformer.feed_forward import FeedForward
+from mflux.models.transformer.joint_attention import JointAttention
 
 
 class JointTransformerBlock(nn.Module):

src/mflux/models/transformer/single_transformer_block.py

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.transformer.ada_layer_norm_zero_single import AdaLayerNormZeroSingle
-from flux_1.models.transformer.single_block_attention import SingleBlockAttention
+from mflux.models.transformer.ada_layer_norm_zero_single import AdaLayerNormZeroSingle
+from mflux.models.transformer.single_block_attention import SingleBlockAttention
 
 
 class SingleTransformerBlock(nn.Module):

src/mflux/models/transformer/time_text_embed.py

@@ -2,11 +2,11 @@ import math
 from mlx import nn
 import mlx.core as mx
 
-from flux_1.config.config import Config
-from flux_1.config.model_config import ModelConfig
-from flux_1.models.transformer.text_embedder import TextEmbedder
-from flux_1.models.transformer.timestep_embedder import TimestepEmbedder
-from flux_1.models.transformer.guidance_embedder import GuidanceEmbedder
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.models.transformer.text_embedder import TextEmbedder
+from mflux.models.transformer.timestep_embedder import TimestepEmbedder
+from mflux.models.transformer.guidance_embedder import GuidanceEmbedder
 
 
 class TimeTextEmbed(nn.Module):

src/mflux/models/transformer/transformer.py

@@ -1,13 +1,13 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.config.model_config import ModelConfig
-from flux_1.config.runtime_config import RuntimeConfig
-from flux_1.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
-from flux_1.models.transformer.embed_nd import EmbedND
-from flux_1.models.transformer.joint_transformer_block import JointTransformerBlock
-from flux_1.models.transformer.single_transformer_block import SingleTransformerBlock
-from flux_1.models.transformer.time_text_embed import TimeTextEmbed
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
+from mflux.models.transformer.embed_nd import EmbedND
+from mflux.models.transformer.joint_transformer_block import JointTransformerBlock
+from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
+from mflux.models.transformer.time_text_embed import TimeTextEmbed
 
 
 class Transformer(nn.Module):

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

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.config.config import Config
+from mflux.config.config import Config
 
 
 class Attention(nn.Module):

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

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.config.config import Config
+from mflux.config.config import Config
 
 
 class ResnetBlock2D(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.attention import Attention
-from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
+from mflux.models.vae.common.attention import Attention
+from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
 
 
 class UnetMidBlock(nn.Module):

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

@@ -1,7 +1,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from flux_1.config.config import Config
+from mflux.config.config import Config
 
 
 class ConvNormOut(nn.Module):

src/mflux/models/vae/decoder/decoder.py

@@ -1,13 +1,13 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.decoder.conv_in import ConvIn
-from flux_1.models.vae.decoder.conv_norm_out import ConvNormOut
-from flux_1.models.vae.decoder.conv_out import ConvOut
-from flux_1.models.vae.common.unet_mid_block import UnetMidBlock
-from flux_1.models.vae.decoder.up_block_1_or_2 import UpBlock1Or2
-from flux_1.models.vae.decoder.up_block_3 import UpBlock3
-from flux_1.models.vae.decoder.up_block_4 import UpBlock4
+from mflux.models.vae.decoder.conv_in import ConvIn
+from mflux.models.vae.decoder.conv_norm_out import ConvNormOut
+from mflux.models.vae.decoder.conv_out import ConvOut
+from mflux.models.vae.common.unet_mid_block import UnetMidBlock
+from mflux.models.vae.decoder.up_block_1_or_2 import UpBlock1Or2
+from mflux.models.vae.decoder.up_block_3 import UpBlock3
+from mflux.models.vae.decoder.up_block_4 import UpBlock4
 
 
 class Decoder(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
-from flux_1.models.vae.decoder.up_sampler import UpSampler
+from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
+from mflux.models.vae.decoder.up_sampler import UpSampler
 
 
 class UpBlock1Or2(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
-from flux_1.models.vae.decoder.up_sampler import UpSampler
+from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
+from mflux.models.vae.decoder.up_sampler import UpSampler
 
 
 class UpBlock3(nn.Module):

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

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
+from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
 
 
 class UpBlock4(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
-from flux_1.models.vae.encoder.down_sampler import DownSampler
+from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
+from mflux.models.vae.encoder.down_sampler import DownSampler
 
 
 class DownBlock1(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
-from flux_1.models.vae.encoder.down_sampler import DownSampler
+from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
+from mflux.models.vae.encoder.down_sampler import DownSampler
 
 
 class DownBlock2(nn.Module):

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

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
-from flux_1.models.vae.encoder.down_sampler import DownSampler
+from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
+from mflux.models.vae.encoder.down_sampler import DownSampler
 
 
 class DownBlock3(nn.Module):

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

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
+from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
 
 
 class DownBlock4(nn.Module):

src/mflux/models/vae/encoder/encoder.py

@@ -1,15 +1,15 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.common.unet_mid_block import UnetMidBlock
+from mflux.models.vae.common.unet_mid_block import UnetMidBlock
 
-from flux_1.models.vae.encoder.conv_in import ConvIn
-from flux_1.models.vae.encoder.conv_norm_out import ConvNormOut
-from flux_1.models.vae.encoder.conv_out import ConvOut
-from flux_1.models.vae.encoder.down_block_1 import DownBlock1
-from flux_1.models.vae.encoder.down_block_2 import DownBlock2
-from flux_1.models.vae.encoder.down_block_3 import DownBlock3
-from flux_1.models.vae.encoder.down_block_4 import DownBlock4
+from mflux.models.vae.encoder.conv_in import ConvIn
+from mflux.models.vae.encoder.conv_norm_out import ConvNormOut
+from mflux.models.vae.encoder.conv_out import ConvOut
+from mflux.models.vae.encoder.down_block_1 import DownBlock1
+from mflux.models.vae.encoder.down_block_2 import DownBlock2
+from mflux.models.vae.encoder.down_block_3 import DownBlock3
+from mflux.models.vae.encoder.down_block_4 import DownBlock4
 
 
 class Encoder(nn.Module):

src/mflux/models/vae/vae.py

@@ -1,8 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1.models.vae.decoder.decoder import Decoder
-from flux_1.models.vae.encoder.encoder import Encoder
+from mflux.models.vae.decoder.decoder import Decoder
+from mflux.models.vae.encoder.encoder import Encoder
 
 
 class VAE(nn.Module):

src/mflux/post_processing/image.py

@@ -6,7 +6,7 @@ import PIL.Image
 import mlx.core as mx
 import piexif
 
-from flux_1.config.model_config import ModelConfig
+from mflux.config.model_config import ModelConfig
 
 log = logging.getLogger(__name__)
 

src/mflux/post_processing/image_util.py

@@ -3,8 +3,8 @@ from PIL import Image
 import mlx.core as mx
 import numpy as np
 
-from flux_1.config.runtime_config import RuntimeConfig
-from flux_1.post_processing.image import Image
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.post_processing.image import Image
 
 
 class ImageUtil:

src/mflux/save.py

@@ -2,10 +2,10 @@ import argparse
 import os
 import sys
 
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from flux_1.flux import Flux1
-from flux_1.config.model_config import ModelConfig
+from mflux.flux.flux import Flux1
+from mflux.config.model_config import ModelConfig
 
 
 def main():

src/mflux/tokenizer/tokenizer_handler.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import transformers
 from huggingface_hub import snapshot_download
 
-from flux_1.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
 
 
 class TokenizerHandler:

src/mflux/weights/lora_util.py

@@ -31,7 +31,7 @@ class LoraUtil:
         if lora_scale < 0.0 or lora_scale > 1.0:
             raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0 - 1.0] ")
 
-        from flux_1.weights.weight_handler import WeightHandler
+        from mflux.weights.weight_handler import WeightHandler
         lora_transformer, _ = WeightHandler.load_transformer(lora_path=lora_file)
         LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale)
 

src/mflux/weights/weight_handler.py

@@ -4,8 +4,8 @@ import mlx.core as mx
 from huggingface_hub import snapshot_download
 from mlx.utils import tree_unflatten
 
-from flux_1.weights.lora_util import LoraUtil
-from flux_1.weights.weight_util import WeightUtil
+from mflux.weights.lora_util import LoraUtil
+from mflux.weights.weight_util import WeightUtil
 
 
 class WeightHandler:

src/mflux/weights/weight_util.py

@@ -1,4 +1,4 @@
-from flux_1.config.config import Config
+from mflux.config.config import Config
 
 
 class WeightUtil: