diff --git a/README.md b/README.md index 315e99a..cf70d94 100644 --- a/README.md +++ b/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) --- diff --git a/main.py b/main.py index 0372c54..ab6a2ad 100644 --- a/main.py +++ b/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(): diff --git a/src/flux_1_schnell/__init__.py b/src/flux_1/__init__.py similarity index 100% rename from src/flux_1_schnell/__init__.py rename to src/flux_1/__init__.py diff --git a/src/flux_1_schnell/config/__init__.py b/src/flux_1/config/__init__.py similarity index 100% rename from src/flux_1_schnell/config/__init__.py rename to src/flux_1/config/__init__.py diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1/config/config.py similarity index 100% rename from src/flux_1_schnell/config/config.py rename to src/flux_1/config/config.py diff --git a/src/flux_1_schnell/config/model_config.py b/src/flux_1/config/model_config.py similarity index 100% rename from src/flux_1_schnell/config/model_config.py rename to src/flux_1/config/model_config.py diff --git a/src/flux_1_schnell/config/runtime_config.py b/src/flux_1/config/runtime_config.py similarity index 94% rename from src/flux_1_schnell/config/runtime_config.py rename to src/flux_1/config/runtime_config.py index e4b3937..e56c9f0 100644 --- a/src/flux_1_schnell/config/runtime_config.py +++ b/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: diff --git a/src/flux_1_schnell/flux.py b/src/flux_1/flux.py similarity index 78% rename from src/flux_1_schnell/flux.py rename to src/flux_1/flux.py index 4b38c2c..8ffbab7 100644 --- a/src/flux_1_schnell/flux.py +++ b/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: diff --git a/src/flux_1_schnell/latent_creator/__init__.py b/src/flux_1/latent_creator/__init__.py similarity index 100% rename from src/flux_1_schnell/latent_creator/__init__.py rename to src/flux_1/latent_creator/__init__.py diff --git a/src/flux_1_schnell/latent_creator/latent_creator.py b/src/flux_1/latent_creator/latent_creator.py similarity index 100% rename from src/flux_1_schnell/latent_creator/latent_creator.py rename to src/flux_1/latent_creator/latent_creator.py diff --git a/src/flux_1_schnell/models/__init__.py b/src/flux_1/models/__init__.py similarity index 100% rename from src/flux_1_schnell/models/__init__.py rename to src/flux_1/models/__init__.py diff --git a/src/flux_1_schnell/models/text_encoder/__init__.py b/src/flux_1/models/text_encoder/__init__.py similarity index 100% rename from src/flux_1_schnell/models/text_encoder/__init__.py rename to src/flux_1/models/text_encoder/__init__.py diff --git a/src/flux_1_schnell/models/text_encoder/clip_encoder/__init__.py b/src/flux_1/models/text_encoder/clip_encoder/__init__.py similarity index 100% rename from src/flux_1_schnell/models/text_encoder/clip_encoder/__init__.py rename to src/flux_1/models/text_encoder/clip_encoder/__init__.py diff --git a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_embeddings.py b/src/flux_1/models/text_encoder/clip_encoder/clip_embeddings.py similarity index 91% rename from src/flux_1_schnell/models/text_encoder/clip_encoder/clip_embeddings.py rename to src/flux_1/models/text_encoder/clip_encoder/clip_embeddings.py index d962f3b..c6de94f 100644 --- a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_embeddings.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_encoder.py b/src/flux_1/models/text_encoder/clip_encoder/clip_encoder.py similarity index 83% rename from src/flux_1_schnell/models/text_encoder/clip_encoder/clip_encoder.py rename to src/flux_1/models/text_encoder/clip_encoder/clip_encoder.py index 4f453e1..204f77a 100644 --- a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_encoder.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_encoder_layer.py b/src/flux_1/models/text_encoder/clip_encoder/clip_encoder_layer.py similarity index 83% rename from src/flux_1_schnell/models/text_encoder/clip_encoder/clip_encoder_layer.py rename to src/flux_1/models/text_encoder/clip_encoder/clip_encoder_layer.py index 01a5a8c..cb7c9f3 100644 --- a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_encoder_layer.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_mlp.py b/src/flux_1/models/text_encoder/clip_encoder/clip_mlp.py similarity index 100% rename from src/flux_1_schnell/models/text_encoder/clip_encoder/clip_mlp.py rename to src/flux_1/models/text_encoder/clip_encoder/clip_mlp.py diff --git a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_sdpa_attention.py b/src/flux_1/models/text_encoder/clip_encoder/clip_sdpa_attention.py similarity index 100% rename from src/flux_1_schnell/models/text_encoder/clip_encoder/clip_sdpa_attention.py rename to src/flux_1/models/text_encoder/clip_encoder/clip_sdpa_attention.py diff --git a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_text_model.py b/src/flux_1/models/text_encoder/clip_encoder/clip_text_model.py similarity index 87% rename from src/flux_1_schnell/models/text_encoder/clip_encoder/clip_text_model.py rename to src/flux_1/models/text_encoder/clip_encoder/clip_text_model.py index f4743e5..5a7e479 100644 --- a/src/flux_1_schnell/models/text_encoder/clip_encoder/clip_text_model.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/clip_encoder/encoder_clip.py b/src/flux_1/models/text_encoder/clip_encoder/encoder_clip.py similarity index 85% rename from src/flux_1_schnell/models/text_encoder/clip_encoder/encoder_clip.py rename to src/flux_1/models/text_encoder/clip_encoder/encoder_clip.py index 5ec12e6..40452a0 100644 --- a/src/flux_1_schnell/models/text_encoder/clip_encoder/encoder_clip.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/__init__.py b/src/flux_1/models/text_encoder/t5_encoder/__init__.py similarity index 100% rename from src/flux_1_schnell/models/text_encoder/t5_encoder/__init__.py rename to src/flux_1/models/text_encoder/t5_encoder/__init__.py diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_attention.py b/src/flux_1/models/text_encoder/t5_encoder/t5_attention.py similarity index 74% rename from src/flux_1_schnell/models/text_encoder/t5_encoder/t5_attention.py rename to src/flux_1/models/text_encoder/t5_encoder/t5_attention.py index 5479898..5ff8da7 100644 --- a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_attention.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_block.py b/src/flux_1/models/text_encoder/t5_encoder/t5_block.py similarity index 72% rename from src/flux_1_schnell/models/text_encoder/t5_encoder/t5_block.py rename to src/flux_1/models/text_encoder/t5_encoder/t5_block.py index a6017ac..481dc40 100644 --- a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_block.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_dense_relu_dense.py b/src/flux_1/models/text_encoder/t5_encoder/t5_dense_relu_dense.py similarity index 100% rename from src/flux_1_schnell/models/text_encoder/t5_encoder/t5_dense_relu_dense.py rename to src/flux_1/models/text_encoder/t5_encoder/t5_dense_relu_dense.py diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_encoder.py b/src/flux_1/models/text_encoder/t5_encoder/t5_encoder.py similarity index 81% rename from src/flux_1_schnell/models/text_encoder/t5_encoder/t5_encoder.py rename to src/flux_1/models/text_encoder/t5_encoder/t5_encoder.py index f1faf27..d893f86 100644 --- a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_encoder.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_feed_forward.py b/src/flux_1/models/text_encoder/t5_encoder/t5_feed_forward.py similarity index 72% rename from src/flux_1_schnell/models/text_encoder/t5_encoder/t5_feed_forward.py rename to src/flux_1/models/text_encoder/t5_encoder/t5_feed_forward.py index 3f7e56f..e3d7792 100644 --- a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_feed_forward.py +++ b/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): diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_layer_norm.py b/src/flux_1/models/text_encoder/t5_encoder/t5_layer_norm.py similarity index 100% rename from src/flux_1_schnell/models/text_encoder/t5_encoder/t5_layer_norm.py rename to src/flux_1/models/text_encoder/t5_encoder/t5_layer_norm.py diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_self_attention.py b/src/flux_1/models/text_encoder/t5_encoder/t5_self_attention.py similarity index 100% rename from src/flux_1_schnell/models/text_encoder/t5_encoder/t5_self_attention.py rename to src/flux_1/models/text_encoder/t5_encoder/t5_self_attention.py diff --git a/src/flux_1_schnell/models/transformer/__init__.py b/src/flux_1/models/transformer/__init__.py similarity index 100% rename from src/flux_1_schnell/models/transformer/__init__.py rename to src/flux_1/models/transformer/__init__.py diff --git a/src/flux_1_schnell/models/transformer/ada_layer_norm_continous.py b/src/flux_1/models/transformer/ada_layer_norm_continous.py similarity index 94% rename from src/flux_1_schnell/models/transformer/ada_layer_norm_continous.py rename to src/flux_1/models/transformer/ada_layer_norm_continous.py index 6731303..eadc253 100644 --- a/src/flux_1_schnell/models/transformer/ada_layer_norm_continous.py +++ b/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): diff --git a/src/flux_1_schnell/models/transformer/ada_layer_norm_zero.py b/src/flux_1/models/transformer/ada_layer_norm_zero.py similarity index 100% rename from src/flux_1_schnell/models/transformer/ada_layer_norm_zero.py rename to src/flux_1/models/transformer/ada_layer_norm_zero.py diff --git a/src/flux_1_schnell/models/transformer/ada_layer_norm_zero_single.py b/src/flux_1/models/transformer/ada_layer_norm_zero_single.py similarity index 100% rename from src/flux_1_schnell/models/transformer/ada_layer_norm_zero_single.py rename to src/flux_1/models/transformer/ada_layer_norm_zero_single.py diff --git a/src/flux_1_schnell/models/transformer/embed_nd.py b/src/flux_1/models/transformer/embed_nd.py similarity index 100% rename from src/flux_1_schnell/models/transformer/embed_nd.py rename to src/flux_1/models/transformer/embed_nd.py diff --git a/src/flux_1_schnell/models/transformer/feed_forward.py b/src/flux_1/models/transformer/feed_forward.py similarity index 100% rename from src/flux_1_schnell/models/transformer/feed_forward.py rename to src/flux_1/models/transformer/feed_forward.py diff --git a/src/flux_1_schnell/models/transformer/guidance_embedder.py b/src/flux_1/models/transformer/guidance_embedder.py similarity index 100% rename from src/flux_1_schnell/models/transformer/guidance_embedder.py rename to src/flux_1/models/transformer/guidance_embedder.py diff --git a/src/flux_1_schnell/models/transformer/joint_attention.py b/src/flux_1/models/transformer/joint_attention.py similarity index 100% rename from src/flux_1_schnell/models/transformer/joint_attention.py rename to src/flux_1/models/transformer/joint_attention.py diff --git a/src/flux_1_schnell/models/transformer/joint_transformer_block.py b/src/flux_1/models/transformer/joint_transformer_block.py similarity index 91% rename from src/flux_1_schnell/models/transformer/joint_transformer_block.py rename to src/flux_1/models/transformer/joint_transformer_block.py index a5b8497..0558099 100644 --- a/src/flux_1_schnell/models/transformer/joint_transformer_block.py +++ b/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): diff --git a/src/flux_1_schnell/models/transformer/single_block_attention.py b/src/flux_1/models/transformer/single_block_attention.py similarity index 100% rename from src/flux_1_schnell/models/transformer/single_block_attention.py rename to src/flux_1/models/transformer/single_block_attention.py diff --git a/src/flux_1_schnell/models/transformer/single_transformer_block.py b/src/flux_1/models/transformer/single_transformer_block.py similarity index 86% rename from src/flux_1_schnell/models/transformer/single_transformer_block.py rename to src/flux_1/models/transformer/single_transformer_block.py index 4d6797b..0c11081 100644 --- a/src/flux_1_schnell/models/transformer/single_transformer_block.py +++ b/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): diff --git a/src/flux_1_schnell/models/transformer/text_embedder.py b/src/flux_1/models/transformer/text_embedder.py similarity index 100% rename from src/flux_1_schnell/models/transformer/text_embedder.py rename to src/flux_1/models/transformer/text_embedder.py diff --git a/src/flux_1_schnell/models/transformer/time_text_embed.py b/src/flux_1/models/transformer/time_text_embed.py similarity index 84% rename from src/flux_1_schnell/models/transformer/time_text_embed.py rename to src/flux_1/models/transformer/time_text_embed.py index 596d271..ed2c335 100644 --- a/src/flux_1_schnell/models/transformer/time_text_embed.py +++ b/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): diff --git a/src/flux_1_schnell/models/transformer/timestep_embedder.py b/src/flux_1/models/transformer/timestep_embedder.py similarity index 100% rename from src/flux_1_schnell/models/transformer/timestep_embedder.py rename to src/flux_1/models/transformer/timestep_embedder.py diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1/models/transformer/transformer.py similarity index 86% rename from src/flux_1_schnell/models/transformer/transformer.py rename to src/flux_1/models/transformer/transformer.py index 15b6119..9caab79 100644 --- a/src/flux_1_schnell/models/transformer/transformer.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/__init__.py b/src/flux_1/models/vae/__init__.py similarity index 100% rename from src/flux_1_schnell/models/vae/__init__.py rename to src/flux_1/models/vae/__init__.py diff --git a/src/flux_1_schnell/models/vae/common/__init__.py b/src/flux_1/models/vae/common/__init__.py similarity index 100% rename from src/flux_1_schnell/models/vae/common/__init__.py rename to src/flux_1/models/vae/common/__init__.py diff --git a/src/flux_1_schnell/models/vae/common/attention.py b/src/flux_1/models/vae/common/attention.py similarity index 95% rename from src/flux_1_schnell/models/vae/common/attention.py rename to src/flux_1/models/vae/common/attention.py index 00c7171..47cfb06 100644 --- a/src/flux_1_schnell/models/vae/common/attention.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/common/resnet_block_2d.py b/src/flux_1/models/vae/common/resnet_block_2d.py similarity index 97% rename from src/flux_1_schnell/models/vae/common/resnet_block_2d.py rename to src/flux_1/models/vae/common/resnet_block_2d.py index da09751..945879e 100644 --- a/src/flux_1_schnell/models/vae/common/resnet_block_2d.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/common/unet_mid_block.py b/src/flux_1/models/vae/common/unet_mid_block.py similarity index 83% rename from src/flux_1_schnell/models/vae/common/unet_mid_block.py rename to src/flux_1/models/vae/common/unet_mid_block.py index e3e943c..946c63b 100644 --- a/src/flux_1_schnell/models/vae/common/unet_mid_block.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/decoder/__init__.py b/src/flux_1/models/vae/decoder/__init__.py similarity index 100% rename from src/flux_1_schnell/models/vae/decoder/__init__.py rename to src/flux_1/models/vae/decoder/__init__.py diff --git a/src/flux_1_schnell/models/vae/decoder/conv_in.py b/src/flux_1/models/vae/decoder/conv_in.py similarity index 100% rename from src/flux_1_schnell/models/vae/decoder/conv_in.py rename to src/flux_1/models/vae/decoder/conv_in.py diff --git a/src/flux_1_schnell/models/vae/decoder/conv_norm_out.py b/src/flux_1/models/vae/decoder/conv_norm_out.py similarity index 92% rename from src/flux_1_schnell/models/vae/decoder/conv_norm_out.py rename to src/flux_1/models/vae/decoder/conv_norm_out.py index 525b6fe..0420bd1 100644 --- a/src/flux_1_schnell/models/vae/decoder/conv_norm_out.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/decoder/conv_out.py b/src/flux_1/models/vae/decoder/conv_out.py similarity index 100% rename from src/flux_1_schnell/models/vae/decoder/conv_out.py rename to src/flux_1/models/vae/decoder/conv_out.py diff --git a/src/flux_1_schnell/models/vae/decoder/decoder.py b/src/flux_1/models/vae/decoder/decoder.py similarity index 63% rename from src/flux_1_schnell/models/vae/decoder/decoder.py rename to src/flux_1/models/vae/decoder/decoder.py index 79b7da8..58eb408 100644 --- a/src/flux_1_schnell/models/vae/decoder/decoder.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/decoder/up_block_1_or_2.py b/src/flux_1/models/vae/decoder/up_block_1_or_2.py similarity index 86% rename from src/flux_1_schnell/models/vae/decoder/up_block_1_or_2.py rename to src/flux_1/models/vae/decoder/up_block_1_or_2.py index 9ff0e54..9ce2ae3 100644 --- a/src/flux_1_schnell/models/vae/decoder/up_block_1_or_2.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/decoder/up_block_3.py b/src/flux_1/models/vae/decoder/up_block_3.py similarity index 87% rename from src/flux_1_schnell/models/vae/decoder/up_block_3.py rename to src/flux_1/models/vae/decoder/up_block_3.py index e6eae56..e25fec4 100644 --- a/src/flux_1_schnell/models/vae/decoder/up_block_3.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/decoder/up_block_4.py b/src/flux_1/models/vae/decoder/up_block_4.py similarity index 91% rename from src/flux_1_schnell/models/vae/decoder/up_block_4.py rename to src/flux_1/models/vae/decoder/up_block_4.py index d1fbd97..b599d54 100644 --- a/src/flux_1_schnell/models/vae/decoder/up_block_4.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/decoder/up_sampler.py b/src/flux_1/models/vae/decoder/up_sampler.py similarity index 100% rename from src/flux_1_schnell/models/vae/decoder/up_sampler.py rename to src/flux_1/models/vae/decoder/up_sampler.py diff --git a/src/flux_1_schnell/models/vae/encoder/__init__.py b/src/flux_1/models/vae/encoder/__init__.py similarity index 100% rename from src/flux_1_schnell/models/vae/encoder/__init__.py rename to src/flux_1/models/vae/encoder/__init__.py diff --git a/src/flux_1_schnell/models/vae/encoder/conv_in.py b/src/flux_1/models/vae/encoder/conv_in.py similarity index 100% rename from src/flux_1_schnell/models/vae/encoder/conv_in.py rename to src/flux_1/models/vae/encoder/conv_in.py diff --git a/src/flux_1_schnell/models/vae/encoder/conv_norm_out.py b/src/flux_1/models/vae/encoder/conv_norm_out.py similarity index 100% rename from src/flux_1_schnell/models/vae/encoder/conv_norm_out.py rename to src/flux_1/models/vae/encoder/conv_norm_out.py diff --git a/src/flux_1_schnell/models/vae/encoder/conv_out.py b/src/flux_1/models/vae/encoder/conv_out.py similarity index 100% rename from src/flux_1_schnell/models/vae/encoder/conv_out.py rename to src/flux_1/models/vae/encoder/conv_out.py diff --git a/src/flux_1_schnell/models/vae/encoder/down_block_1.py b/src/flux_1/models/vae/encoder/down_block_1.py similarity index 84% rename from src/flux_1_schnell/models/vae/encoder/down_block_1.py rename to src/flux_1/models/vae/encoder/down_block_1.py index 4f107d0..09d3c4a 100644 --- a/src/flux_1_schnell/models/vae/encoder/down_block_1.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/encoder/down_block_2.py b/src/flux_1/models/vae/encoder/down_block_2.py similarity index 85% rename from src/flux_1_schnell/models/vae/encoder/down_block_2.py rename to src/flux_1/models/vae/encoder/down_block_2.py index 90c1b44..fb261c6 100644 --- a/src/flux_1_schnell/models/vae/encoder/down_block_2.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/encoder/down_block_3.py b/src/flux_1/models/vae/encoder/down_block_3.py similarity index 85% rename from src/flux_1_schnell/models/vae/encoder/down_block_3.py rename to src/flux_1/models/vae/encoder/down_block_3.py index 87de0bb..86eddb9 100644 --- a/src/flux_1_schnell/models/vae/encoder/down_block_3.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/encoder/down_block_4.py b/src/flux_1/models/vae/encoder/down_block_4.py similarity index 88% rename from src/flux_1_schnell/models/vae/encoder/down_block_4.py rename to src/flux_1/models/vae/encoder/down_block_4.py index cd11669..b8eb214 100644 --- a/src/flux_1_schnell/models/vae/encoder/down_block_4.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/encoder/down_sampler.py b/src/flux_1/models/vae/encoder/down_sampler.py similarity index 100% rename from src/flux_1_schnell/models/vae/encoder/down_sampler.py rename to src/flux_1/models/vae/encoder/down_sampler.py diff --git a/src/flux_1_schnell/models/vae/encoder/encoder.py b/src/flux_1/models/vae/encoder/encoder.py similarity index 59% rename from src/flux_1_schnell/models/vae/encoder/encoder.py rename to src/flux_1/models/vae/encoder/encoder.py index c401319..14d3898 100644 --- a/src/flux_1_schnell/models/vae/encoder/encoder.py +++ b/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): diff --git a/src/flux_1_schnell/models/vae/vae.py b/src/flux_1/models/vae/vae.py similarity index 85% rename from src/flux_1_schnell/models/vae/vae.py rename to src/flux_1/models/vae/vae.py index 1aad0f4..e6de9e2 100644 --- a/src/flux_1_schnell/models/vae/vae.py +++ b/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): diff --git a/src/flux_1_schnell/post_processing/__init__.py b/src/flux_1/post_processing/__init__.py similarity index 100% rename from src/flux_1_schnell/post_processing/__init__.py rename to src/flux_1/post_processing/__init__.py diff --git a/src/flux_1_schnell/post_processing/image_util.py b/src/flux_1/post_processing/image_util.py similarity index 100% rename from src/flux_1_schnell/post_processing/image_util.py rename to src/flux_1/post_processing/image_util.py diff --git a/src/flux_1_schnell/tokenizer/__init__.py b/src/flux_1/tokenizer/__init__.py similarity index 100% rename from src/flux_1_schnell/tokenizer/__init__.py rename to src/flux_1/tokenizer/__init__.py diff --git a/src/flux_1_schnell/tokenizer/clip_tokenizer.py b/src/flux_1/tokenizer/clip_tokenizer.py similarity index 100% rename from src/flux_1_schnell/tokenizer/clip_tokenizer.py rename to src/flux_1/tokenizer/clip_tokenizer.py diff --git a/src/flux_1_schnell/tokenizer/t5_tokenizer.py b/src/flux_1/tokenizer/t5_tokenizer.py similarity index 100% rename from src/flux_1_schnell/tokenizer/t5_tokenizer.py rename to src/flux_1/tokenizer/t5_tokenizer.py diff --git a/src/flux_1_schnell/tokenizer/tokenizer_handler.py b/src/flux_1/tokenizer/tokenizer_handler.py similarity index 90% rename from src/flux_1_schnell/tokenizer/tokenizer_handler.py rename to src/flux_1/tokenizer/tokenizer_handler.py index 0269e53..58665c6 100644 --- a/src/flux_1_schnell/tokenizer/tokenizer_handler.py +++ b/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: diff --git a/src/flux_1_schnell/weights/__init__.py b/src/flux_1/weights/__init__.py similarity index 100% rename from src/flux_1_schnell/weights/__init__.py rename to src/flux_1/weights/__init__.py diff --git a/src/flux_1_schnell/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py similarity index 99% rename from src/flux_1_schnell/weights/weight_handler.py rename to src/flux_1/weights/weight_handler.py index e9147c2..21eed11 100644 --- a/src/flux_1_schnell/weights/weight_handler.py +++ b/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: diff --git a/src/flux_1_schnell/assets/comparison1.jpg b/src/flux_1_schnell/assets/comparison1.jpg deleted file mode 100644 index 8b36c66..0000000 Binary files a/src/flux_1_schnell/assets/comparison1.jpg and /dev/null differ diff --git a/src/flux_1_schnell/assets/comparison2.jpg b/src/flux_1_schnell/assets/comparison2.jpg deleted file mode 100644 index caab80d..0000000 Binary files a/src/flux_1_schnell/assets/comparison2.jpg and /dev/null differ diff --git a/src/flux_1_schnell/assets/comparison3.jpg b/src/flux_1_schnell/assets/comparison3.jpg deleted file mode 100644 index f73a400..0000000 Binary files a/src/flux_1_schnell/assets/comparison3.jpg and /dev/null differ diff --git a/src/flux_1_schnell/assets/comparison4.jpg b/src/flux_1_schnell/assets/comparison4.jpg deleted file mode 100644 index 8342492..0000000 Binary files a/src/flux_1_schnell/assets/comparison4.jpg and /dev/null differ diff --git a/src/flux_1_schnell/assets/comparison5.jpg b/src/flux_1_schnell/assets/comparison5.jpg deleted file mode 100644 index d9bc3da..0000000 Binary files a/src/flux_1_schnell/assets/comparison5.jpg and /dev/null differ diff --git a/src/flux_1_schnell/assets/logo.png b/src/flux_1_schnell/assets/logo.png deleted file mode 100644 index f791a52..0000000 Binary files a/src/flux_1_schnell/assets/logo.png and /dev/null differ