Rename flux_1_schnell folder to flux_1

README.md

@@ -1,5 +1,5 @@
 
-![image](src/flux_1_schnell/assets/logo.png)
+![image](src/flux_1/assets/logo.png)
 *A MLX port of FLUX based on the Huggingface Diffusers implementation.*
 
 ### About
@@ -74,9 +74,9 @@ import sys
 
 sys.path.append("/path/to/mflux/src")
 
-from flux_1_schnell.config.config import Config
-from flux_1_schnell.flux import Flux1
-from flux_1_schnell.post_processing.image_util import ImageUtil
+from flux_1.config.config import Config
+from flux_1.flux import Flux1
+from flux_1.post_processing.image_util import ImageUtil
 
 flux = Flux1.from_repo("black-forest-labs/FLUX.1-schnell")
 
@@ -127,27 +127,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_schnell/config/config.py) class. 
+The precision for MFLUX can be set in the [Config](src/flux_1/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)
+![image](src/flux_1/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)
+![image](src/flux_1/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)
+![image](src/flux_1/assets/comparison3.jpg)
 
 ---
 
@@ -155,14 +155,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_schnell/assets/comparison4.jpg)
+![image](src/flux_1/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)
+![image](src/flux_1/assets/comparison5.jpg)
 
 ---
 

main.py

@@ -5,9 +5,9 @@ import time
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
 
-from flux_1_schnell.config.config import Config
-from flux_1_schnell.flux import Flux1
-from flux_1_schnell.post_processing.image_util import ImageUtil
+from flux_1.config.config import Config
+from flux_1.flux import Flux1
+from flux_1.post_processing.image_util import ImageUtil
 
 
 def main():

src/flux_1/config/runtime_config.py

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

src/flux_1/flux.py

@@ -3,19 +3,19 @@ import mlx.core as mx
 from PIL import Image
 from tqdm import tqdm
 
-from flux_1_schnell.config.config import Config
-from flux_1_schnell.config.runtime_config import RuntimeConfig
-from flux_1_schnell.latent_creator.latent_creator import LatentCreator
-from flux_1_schnell.config.model_config import ModelConfig
-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.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.tokenizer.clip_tokenizer import TokenizerCLIP
-from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
-from flux_1_schnell.tokenizer.tokenizer_handler import TokenizerHandler
-from flux_1_schnell.weights.weight_handler import WeightHandler
+from flux_1.config.config import Config
+from flux_1.config.runtime_config import RuntimeConfig
+from flux_1.latent_creator.latent_creator import LatentCreator
+from flux_1.config.model_config import ModelConfig
+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_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
 
 
 class Flux1:

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

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

src/flux_1/models/text_encoder/clip_encoder/clip_encoder.py

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

src/flux_1/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_schnell.models.text_encoder.clip_encoder.clip_mlp import CLIPMLP
-from flux_1_schnell.models.text_encoder.clip_encoder.clip_sdpa_attention import CLIPSdpaAttention
+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
 
 
 class CLIPEncoderLayer(nn.Module):

src/flux_1/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_schnell.models.text_encoder.clip_encoder.clip_embeddings import CLIPEmbeddings
-from flux_1_schnell.models.text_encoder.clip_encoder.encoder_clip import EncoderCLIP
+from flux_1.models.text_encoder.clip_encoder.clip_embeddings import CLIPEmbeddings
+from flux_1.models.text_encoder.clip_encoder.encoder_clip import EncoderCLIP
 
 
 class CLIPTextModel(nn.Module):

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

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

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

@@ -1,8 +1,8 @@
 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
+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
 
 
 class T5Attention(nn.Module):

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

@@ -1,8 +1,8 @@
 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
+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
 
 
 class T5Block(nn.Module):

src/flux_1/models/text_encoder/t5_encoder/t5_encoder.py

@@ -1,8 +1,8 @@
 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
+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
 
 
 class T5Encoder(nn.Module):

src/flux_1/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_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
+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
 
 
 class T5FeedForward(nn.Module):

src/flux_1/models/transformer/ada_layer_norm_continous.py

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

src/flux_1/models/transformer/joint_transformer_block.py

@@ -1,9 +1,9 @@
 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
+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
 
 
 class JointTransformerBlock(nn.Module):

src/flux_1/models/transformer/single_transformer_block.py

@@ -1,8 +1,8 @@
 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
+from flux_1.models.transformer.ada_layer_norm_zero_single import AdaLayerNormZeroSingle
+from flux_1.models.transformer.single_block_attention import SingleBlockAttention
 
 
 class SingleTransformerBlock(nn.Module):

src/flux_1/models/transformer/time_text_embed.py

@@ -2,10 +2,10 @@ 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
-from flux_1_schnell.models.transformer.guidance_embedder import GuidanceEmbedder
+from flux_1.config.config import Config
+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
 
 
 class TimeTextEmbed(nn.Module):

src/flux_1/models/transformer/transformer.py

@@ -1,13 +1,13 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1_schnell.config.config import Config
-from flux_1_schnell.config.runtime_config import RuntimeConfig
-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
+from flux_1.config.config import Config
+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
 
 
 class Transformer(nn.Module):

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

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

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

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

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

@@ -1,8 +1,8 @@
 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
+from flux_1.models.vae.common.attention import Attention
+from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
 
 
 class UnetMidBlock(nn.Module):

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

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

src/flux_1/models/vae/decoder/decoder.py

@@ -1,13 +1,13 @@
 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
+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
 
 
 class Decoder(nn.Module):

src/flux_1/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_schnell.models.vae.common.resnet_block_2d import ResnetBlock2D
-from flux_1_schnell.models.vae.decoder.up_sampler import UpSampler
+from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1.models.vae.decoder.up_sampler import UpSampler
 
 
 class UpBlock1Or2(nn.Module):

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

@@ -1,8 +1,8 @@
 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
+from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1.models.vae.decoder.up_sampler import UpSampler
 
 
 class UpBlock3(nn.Module):

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

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

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

@@ -1,8 +1,8 @@
 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
+from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1.models.vae.encoder.down_sampler import DownSampler
 
 
 class DownBlock1(nn.Module):

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

@@ -1,8 +1,8 @@
 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
+from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1.models.vae.encoder.down_sampler import DownSampler
 
 
 class DownBlock2(nn.Module):

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

@@ -1,8 +1,8 @@
 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
+from flux_1.models.vae.common.resnet_block_2d import ResnetBlock2D
+from flux_1.models.vae.encoder.down_sampler import DownSampler
 
 
 class DownBlock3(nn.Module):

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

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

src/flux_1/models/vae/encoder/encoder.py

@@ -1,15 +1,15 @@
 import mlx.core as mx
 from mlx import nn
 
-from flux_1_schnell.models.vae.common.unet_mid_block import UnetMidBlock
+from flux_1.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
+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
 
 
 class Encoder(nn.Module):

src/flux_1/models/vae/vae.py

@@ -1,8 +1,8 @@
 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
+from flux_1.models.vae.decoder.decoder import Decoder
+from flux_1.models.vae.encoder.encoder import Encoder
 
 
 class VAE(nn.Module):

src/flux_1/tokenizer/tokenizer_handler.py

@@ -3,8 +3,8 @@ from pathlib import Path
 import transformers
 from huggingface_hub import snapshot_download
 
-from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
-from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
+from flux_1.tokenizer.clip_tokenizer import TokenizerCLIP
+from flux_1.tokenizer.t5_tokenizer import TokenizerT5
 
 
 class TokenizerHandler:

src/flux_1/weights/weight_handler.py

@@ -4,7 +4,7 @@ import mlx.core as mx
 from huggingface_hub import snapshot_download
 from mlx.utils import tree_unflatten
 
-from flux_1_schnell.config.config import Config
+from flux_1.config.config import Config
 
 
 class WeightHandler: