From d2b38011adffba8980c1c9571d99b1792a272958 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Fri, 6 Sep 2024 21:06:41 +0200 Subject: [PATCH] Rename flux_1 to mflux --- README.md | 26 +++++++++--------- main.py | 6 ++-- save.py | 4 +-- src/{flux_1 => mflux}/__init__.py | 0 src/{flux_1 => mflux}/assets/comparison1.jpg | Bin src/{flux_1 => mflux}/assets/comparison2.jpg | Bin src/{flux_1 => mflux}/assets/comparison3.jpg | Bin src/{flux_1 => mflux}/assets/comparison4.jpg | Bin src/{flux_1 => mflux}/assets/comparison5.jpg | Bin src/{flux_1 => mflux}/assets/comparison6.jpg | Bin src/{flux_1 => mflux}/assets/logo.png | Bin src/{flux_1 => mflux}/assets/lora1.jpg | Bin src/{flux_1 => mflux}/assets/lora2.jpg | Bin src/{flux_1 => mflux}/assets/lora3.jpg | Bin src/{flux_1 => mflux}/config/__init__.py | 0 src/{flux_1 => mflux}/config/config.py | 0 src/{flux_1 => mflux}/config/model_config.py | 0 .../config/runtime_config.py | 4 +-- src/{flux_1 => mflux}/flux.py | 26 +++++++++--------- src/{flux_1 => mflux}/models/__init__.py | 0 .../models/text_encoder/__init__.py | 0 .../text_encoder/clip_encoder/__init__.py | 0 .../clip_encoder/clip_embeddings.py | 2 +- .../text_encoder/clip_encoder/clip_encoder.py | 2 +- .../clip_encoder/clip_encoder_layer.py | 4 +-- .../text_encoder/clip_encoder/clip_mlp.py | 0 .../clip_encoder/clip_sdpa_attention.py | 0 .../clip_encoder/clip_text_model.py | 4 +-- .../text_encoder/clip_encoder/encoder_clip.py | 2 +- .../text_encoder/t5_encoder/__init__.py | 0 .../text_encoder/t5_encoder/t5_attention.py | 4 +-- .../text_encoder/t5_encoder/t5_block.py | 4 +-- .../t5_encoder/t5_dense_relu_dense.py | 0 .../text_encoder/t5_encoder/t5_encoder.py | 4 +-- .../t5_encoder/t5_feed_forward.py | 4 +-- .../text_encoder/t5_encoder/t5_layer_norm.py | 0 .../t5_encoder/t5_self_attention.py | 0 .../models/transformer/__init__.py | 0 .../transformer/ada_layer_norm_continous.py | 2 +- .../models/transformer/ada_layer_norm_zero.py | 0 .../transformer/ada_layer_norm_zero_single.py | 0 .../models/transformer/embed_nd.py | 0 .../models/transformer/feed_forward.py | 0 .../models/transformer/guidance_embedder.py | 0 .../models/transformer/joint_attention.py | 0 .../transformer/joint_transformer_block.py | 6 ++-- .../transformer/single_block_attention.py | 0 .../transformer/single_transformer_block.py | 4 +-- .../models/transformer/text_embedder.py | 0 .../models/transformer/time_text_embed.py | 10 +++---- .../models/transformer/timestep_embedder.py | 0 .../models/transformer/transformer.py | 14 +++++----- src/{flux_1 => mflux}/models/vae/__init__.py | 0 .../models/vae/common/__init__.py | 0 .../models/vae/common/attention.py | 2 +- .../models/vae/common/resnet_block_2d.py | 2 +- .../models/vae/common/unet_mid_block.py | 4 +-- .../models/vae/decoder/__init__.py | 0 .../models/vae/decoder/conv_in.py | 0 .../models/vae/decoder/conv_norm_out.py | 2 +- .../models/vae/decoder/conv_out.py | 0 .../models/vae/decoder/decoder.py | 14 +++++----- .../models/vae/decoder/up_block_1_or_2.py | 4 +-- .../models/vae/decoder/up_block_3.py | 4 +-- .../models/vae/decoder/up_block_4.py | 2 +- .../models/vae/decoder/up_sampler.py | 0 .../models/vae/encoder/__init__.py | 0 .../models/vae/encoder/conv_in.py | 0 .../models/vae/encoder/conv_norm_out.py | 0 .../models/vae/encoder/conv_out.py | 0 .../models/vae/encoder/down_block_1.py | 4 +-- .../models/vae/encoder/down_block_2.py | 4 +-- .../models/vae/encoder/down_block_3.py | 4 +-- .../models/vae/encoder/down_block_4.py | 2 +- .../models/vae/encoder/down_sampler.py | 0 .../models/vae/encoder/encoder.py | 16 +++++------ src/{flux_1 => mflux}/models/vae/vae.py | 4 +-- .../post_processing/__init__.py | 0 .../post_processing/image.py | 2 +- .../post_processing/image_util.py | 4 +-- src/{flux_1 => mflux}/tokenizer/__init__.py | 0 .../tokenizer/clip_tokenizer.py | 0 .../tokenizer/t5_tokenizer.py | 0 .../tokenizer/tokenizer_handler.py | 2 +- src/{flux_1 => mflux}/weights/__init__.py | 0 src/{flux_1 => mflux}/weights/lora_util.py | 2 +- .../weights/weight_handler.py | 4 +-- src/{flux_1 => mflux}/weights/weight_util.py | 2 +- 88 files changed, 108 insertions(+), 108 deletions(-) rename src/{flux_1 => mflux}/__init__.py (100%) rename src/{flux_1 => mflux}/assets/comparison1.jpg (100%) rename src/{flux_1 => mflux}/assets/comparison2.jpg (100%) rename src/{flux_1 => mflux}/assets/comparison3.jpg (100%) rename src/{flux_1 => mflux}/assets/comparison4.jpg (100%) rename src/{flux_1 => mflux}/assets/comparison5.jpg (100%) rename src/{flux_1 => mflux}/assets/comparison6.jpg (100%) rename src/{flux_1 => mflux}/assets/logo.png (100%) rename src/{flux_1 => mflux}/assets/lora1.jpg (100%) rename src/{flux_1 => mflux}/assets/lora2.jpg (100%) rename src/{flux_1 => mflux}/assets/lora3.jpg (100%) rename src/{flux_1 => mflux}/config/__init__.py (100%) rename src/{flux_1 => mflux}/config/config.py (100%) rename src/{flux_1 => mflux}/config/model_config.py (100%) rename src/{flux_1 => mflux}/config/runtime_config.py (95%) rename src/{flux_1 => mflux}/flux.py (90%) rename src/{flux_1 => mflux}/models/__init__.py (100%) rename src/{flux_1 => mflux}/models/text_encoder/__init__.py (100%) rename src/{flux_1 => mflux}/models/text_encoder/clip_encoder/__init__.py (100%) rename src/{flux_1 => mflux}/models/text_encoder/clip_encoder/clip_embeddings.py (92%) rename src/{flux_1 => mflux}/models/text_encoder/clip_encoder/clip_encoder.py (80%) rename src/{flux_1 => mflux}/models/text_encoder/clip_encoder/clip_encoder_layer.py (84%) rename src/{flux_1 => mflux}/models/text_encoder/clip_encoder/clip_mlp.py (100%) rename src/{flux_1 => mflux}/models/text_encoder/clip_encoder/clip_sdpa_attention.py (100%) rename src/{flux_1 => mflux}/models/text_encoder/clip_encoder/clip_text_model.py (88%) rename src/{flux_1 => mflux}/models/text_encoder/clip_encoder/encoder_clip.py (86%) rename src/{flux_1 => mflux}/models/text_encoder/t5_encoder/__init__.py (100%) rename src/{flux_1 => mflux}/models/text_encoder/t5_encoder/t5_attention.py (75%) rename src/{flux_1 => mflux}/models/text_encoder/t5_encoder/t5_block.py (73%) rename src/{flux_1 => mflux}/models/text_encoder/t5_encoder/t5_dense_relu_dense.py (100%) rename src/{flux_1 => mflux}/models/text_encoder/t5_encoder/t5_encoder.py (80%) rename src/{flux_1 => mflux}/models/text_encoder/t5_encoder/t5_feed_forward.py (74%) rename src/{flux_1 => mflux}/models/text_encoder/t5_encoder/t5_layer_norm.py (100%) rename src/{flux_1 => mflux}/models/text_encoder/t5_encoder/t5_self_attention.py (100%) rename src/{flux_1 => mflux}/models/transformer/__init__.py (100%) rename src/{flux_1 => mflux}/models/transformer/ada_layer_norm_continous.py (95%) rename src/{flux_1 => mflux}/models/transformer/ada_layer_norm_zero.py (100%) rename src/{flux_1 => mflux}/models/transformer/ada_layer_norm_zero_single.py (100%) rename src/{flux_1 => mflux}/models/transformer/embed_nd.py (100%) rename src/{flux_1 => mflux}/models/transformer/feed_forward.py (100%) rename src/{flux_1 => mflux}/models/transformer/guidance_embedder.py (100%) rename src/{flux_1 => mflux}/models/transformer/joint_attention.py (100%) rename src/{flux_1 => mflux}/models/transformer/joint_transformer_block.py (92%) rename src/{flux_1 => mflux}/models/transformer/single_block_attention.py (100%) rename src/{flux_1 => mflux}/models/transformer/single_transformer_block.py (87%) rename src/{flux_1 => mflux}/models/transformer/text_embedder.py (100%) rename src/{flux_1 => mflux}/models/transformer/time_text_embed.py (82%) rename src/{flux_1 => mflux}/models/transformer/timestep_embedder.py (100%) rename src/{flux_1 => mflux}/models/transformer/transformer.py (87%) rename src/{flux_1 => mflux}/models/vae/__init__.py (100%) rename src/{flux_1 => mflux}/models/vae/common/__init__.py (100%) rename src/{flux_1 => mflux}/models/vae/common/attention.py (96%) rename src/{flux_1 => mflux}/models/vae/common/resnet_block_2d.py (98%) rename src/{flux_1 => mflux}/models/vae/common/unet_mid_block.py (84%) rename src/{flux_1 => mflux}/models/vae/decoder/__init__.py (100%) rename src/{flux_1 => mflux}/models/vae/decoder/conv_in.py (100%) rename src/{flux_1 => mflux}/models/vae/decoder/conv_norm_out.py (93%) rename src/{flux_1 => mflux}/models/vae/decoder/conv_out.py (100%) rename src/{flux_1 => mflux}/models/vae/decoder/decoder.py (66%) rename src/{flux_1 => mflux}/models/vae/decoder/up_block_1_or_2.py (88%) rename src/{flux_1 => mflux}/models/vae/decoder/up_block_3.py (88%) rename src/{flux_1 => mflux}/models/vae/decoder/up_block_4.py (92%) rename src/{flux_1 => mflux}/models/vae/decoder/up_sampler.py (100%) rename src/{flux_1 => mflux}/models/vae/encoder/__init__.py (100%) rename src/{flux_1 => mflux}/models/vae/encoder/conv_in.py (100%) rename src/{flux_1 => mflux}/models/vae/encoder/conv_norm_out.py (100%) rename src/{flux_1 => mflux}/models/vae/encoder/conv_out.py (100%) rename src/{flux_1 => mflux}/models/vae/encoder/down_block_1.py (85%) rename src/{flux_1 => mflux}/models/vae/encoder/down_block_2.py (86%) rename src/{flux_1 => mflux}/models/vae/encoder/down_block_3.py (86%) rename src/{flux_1 => mflux}/models/vae/encoder/down_block_4.py (89%) rename src/{flux_1 => mflux}/models/vae/encoder/down_sampler.py (100%) rename src/{flux_1 => mflux}/models/vae/encoder/encoder.py (62%) rename src/{flux_1 => mflux}/models/vae/vae.py (85%) rename src/{flux_1 => mflux}/post_processing/__init__.py (100%) rename src/{flux_1 => mflux}/post_processing/image.py (98%) rename src/{flux_1 => mflux}/post_processing/image_util.py (95%) rename src/{flux_1 => mflux}/tokenizer/__init__.py (100%) rename src/{flux_1 => mflux}/tokenizer/clip_tokenizer.py (100%) rename src/{flux_1 => mflux}/tokenizer/t5_tokenizer.py (100%) rename src/{flux_1 => mflux}/tokenizer/tokenizer_handler.py (94%) rename src/{flux_1 => mflux}/weights/__init__.py (100%) rename src/{flux_1 => mflux}/weights/lora_util.py (98%) rename src/{flux_1 => mflux}/weights/weight_handler.py (98%) rename src/{flux_1 => mflux}/weights/weight_util.py (90%) diff --git a/README.md b/README.md index ccff1fc..07700f3 100644 --- a/README.md +++ b/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 @@ -96,8 +96,8 @@ import sys sys.path.append("/path/to/mflux/src") -from flux_1.config.config import Config -from flux_1.flux import Flux1 +from mflux.config.config import Config +from mflux.flux import Flux1 # Load the model flux = Flux1.from_alias(alias="schnell") # "schnell" or "dev" @@ -157,27 +157,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 +185,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) --- @@ -211,7 +211,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.* @@ -320,7 +320,7 @@ The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) L python main.py --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): @@ -328,7 +328,7 @@ The following example is [Flux_1_Dev_LoRA_Paper-Cutout-Style](https://huggingfac ``` 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" ``` -![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"**.* @@ -348,7 +348,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. diff --git a/main.py b/main.py index 785a5f2..facae1c 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.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 import Flux1 def main(): diff --git a/save.py b/save.py index 482de17..d38f19c 100644 --- a/save.py +++ b/save.py @@ -4,8 +4,8 @@ import sys sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src'))) -from flux_1.flux import Flux1 -from flux_1.config.model_config import ModelConfig +from mflux.flux import Flux1 +from mflux.config.model_config import ModelConfig def main(): diff --git a/src/flux_1/__init__.py b/src/mflux/__init__.py similarity index 100% rename from src/flux_1/__init__.py rename to src/mflux/__init__.py diff --git a/src/flux_1/assets/comparison1.jpg b/src/mflux/assets/comparison1.jpg similarity index 100% rename from src/flux_1/assets/comparison1.jpg rename to src/mflux/assets/comparison1.jpg diff --git a/src/flux_1/assets/comparison2.jpg b/src/mflux/assets/comparison2.jpg similarity index 100% rename from src/flux_1/assets/comparison2.jpg rename to src/mflux/assets/comparison2.jpg diff --git a/src/flux_1/assets/comparison3.jpg b/src/mflux/assets/comparison3.jpg similarity index 100% rename from src/flux_1/assets/comparison3.jpg rename to src/mflux/assets/comparison3.jpg diff --git a/src/flux_1/assets/comparison4.jpg b/src/mflux/assets/comparison4.jpg similarity index 100% rename from src/flux_1/assets/comparison4.jpg rename to src/mflux/assets/comparison4.jpg diff --git a/src/flux_1/assets/comparison5.jpg b/src/mflux/assets/comparison5.jpg similarity index 100% rename from src/flux_1/assets/comparison5.jpg rename to src/mflux/assets/comparison5.jpg diff --git a/src/flux_1/assets/comparison6.jpg b/src/mflux/assets/comparison6.jpg similarity index 100% rename from src/flux_1/assets/comparison6.jpg rename to src/mflux/assets/comparison6.jpg diff --git a/src/flux_1/assets/logo.png b/src/mflux/assets/logo.png similarity index 100% rename from src/flux_1/assets/logo.png rename to src/mflux/assets/logo.png diff --git a/src/flux_1/assets/lora1.jpg b/src/mflux/assets/lora1.jpg similarity index 100% rename from src/flux_1/assets/lora1.jpg rename to src/mflux/assets/lora1.jpg diff --git a/src/flux_1/assets/lora2.jpg b/src/mflux/assets/lora2.jpg similarity index 100% rename from src/flux_1/assets/lora2.jpg rename to src/mflux/assets/lora2.jpg diff --git a/src/flux_1/assets/lora3.jpg b/src/mflux/assets/lora3.jpg similarity index 100% rename from src/flux_1/assets/lora3.jpg rename to src/mflux/assets/lora3.jpg diff --git a/src/flux_1/config/__init__.py b/src/mflux/config/__init__.py similarity index 100% rename from src/flux_1/config/__init__.py rename to src/mflux/config/__init__.py diff --git a/src/flux_1/config/config.py b/src/mflux/config/config.py similarity index 100% rename from src/flux_1/config/config.py rename to src/mflux/config/config.py diff --git a/src/flux_1/config/model_config.py b/src/mflux/config/model_config.py similarity index 100% rename from src/flux_1/config/model_config.py rename to src/mflux/config/model_config.py diff --git a/src/flux_1/config/runtime_config.py b/src/mflux/config/runtime_config.py similarity index 95% rename from src/flux_1/config/runtime_config.py rename to src/mflux/config/runtime_config.py index 6e2e461..33cd634 100644 --- a/src/flux_1/config/runtime_config.py +++ b/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: diff --git a/src/flux_1/flux.py b/src/mflux/flux.py similarity index 90% rename from src/flux_1/flux.py rename to src/mflux/flux.py index 50511f2..35d9356 100644 --- a/src/flux_1/flux.py +++ b/src/mflux/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: diff --git a/src/flux_1/models/__init__.py b/src/mflux/models/__init__.py similarity index 100% rename from src/flux_1/models/__init__.py rename to src/mflux/models/__init__.py diff --git a/src/flux_1/models/text_encoder/__init__.py b/src/mflux/models/text_encoder/__init__.py similarity index 100% rename from src/flux_1/models/text_encoder/__init__.py rename to src/mflux/models/text_encoder/__init__.py diff --git a/src/flux_1/models/text_encoder/clip_encoder/__init__.py b/src/mflux/models/text_encoder/clip_encoder/__init__.py similarity index 100% rename from src/flux_1/models/text_encoder/clip_encoder/__init__.py rename to src/mflux/models/text_encoder/clip_encoder/__init__.py diff --git a/src/flux_1/models/text_encoder/clip_encoder/clip_embeddings.py b/src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py similarity index 92% rename from src/flux_1/models/text_encoder/clip_encoder/clip_embeddings.py rename to src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py index c6de94f..9c4fd83 100644 --- a/src/flux_1/models/text_encoder/clip_encoder/clip_embeddings.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/clip_encoder/clip_encoder.py b/src/mflux/models/text_encoder/clip_encoder/clip_encoder.py similarity index 80% rename from src/flux_1/models/text_encoder/clip_encoder/clip_encoder.py rename to src/mflux/models/text_encoder/clip_encoder/clip_encoder.py index 8a3a509..a8ab552 100644 --- a/src/flux_1/models/text_encoder/clip_encoder/clip_encoder.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/clip_encoder/clip_encoder_layer.py b/src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py similarity index 84% rename from src/flux_1/models/text_encoder/clip_encoder/clip_encoder_layer.py rename to src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py index cb7c9f3..3d20300 100644 --- a/src/flux_1/models/text_encoder/clip_encoder/clip_encoder_layer.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/clip_encoder/clip_mlp.py b/src/mflux/models/text_encoder/clip_encoder/clip_mlp.py similarity index 100% rename from src/flux_1/models/text_encoder/clip_encoder/clip_mlp.py rename to src/mflux/models/text_encoder/clip_encoder/clip_mlp.py diff --git a/src/flux_1/models/text_encoder/clip_encoder/clip_sdpa_attention.py b/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py similarity index 100% rename from src/flux_1/models/text_encoder/clip_encoder/clip_sdpa_attention.py rename to src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py diff --git a/src/flux_1/models/text_encoder/clip_encoder/clip_text_model.py b/src/mflux/models/text_encoder/clip_encoder/clip_text_model.py similarity index 88% rename from src/flux_1/models/text_encoder/clip_encoder/clip_text_model.py rename to src/mflux/models/text_encoder/clip_encoder/clip_text_model.py index 5a7e479..1f25fa6 100644 --- a/src/flux_1/models/text_encoder/clip_encoder/clip_text_model.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/clip_encoder/encoder_clip.py b/src/mflux/models/text_encoder/clip_encoder/encoder_clip.py similarity index 86% rename from src/flux_1/models/text_encoder/clip_encoder/encoder_clip.py rename to src/mflux/models/text_encoder/clip_encoder/encoder_clip.py index 40452a0..7ec2b12 100644 --- a/src/flux_1/models/text_encoder/clip_encoder/encoder_clip.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/t5_encoder/__init__.py b/src/mflux/models/text_encoder/t5_encoder/__init__.py similarity index 100% rename from src/flux_1/models/text_encoder/t5_encoder/__init__.py rename to src/mflux/models/text_encoder/t5_encoder/__init__.py diff --git a/src/flux_1/models/text_encoder/t5_encoder/t5_attention.py b/src/mflux/models/text_encoder/t5_encoder/t5_attention.py similarity index 75% rename from src/flux_1/models/text_encoder/t5_encoder/t5_attention.py rename to src/mflux/models/text_encoder/t5_encoder/t5_attention.py index 5ff8da7..f7ba4aa 100644 --- a/src/flux_1/models/text_encoder/t5_encoder/t5_attention.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/t5_encoder/t5_block.py b/src/mflux/models/text_encoder/t5_encoder/t5_block.py similarity index 73% rename from src/flux_1/models/text_encoder/t5_encoder/t5_block.py rename to src/mflux/models/text_encoder/t5_encoder/t5_block.py index 481dc40..48f0fdc 100644 --- a/src/flux_1/models/text_encoder/t5_encoder/t5_block.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/t5_encoder/t5_dense_relu_dense.py b/src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py similarity index 100% rename from src/flux_1/models/text_encoder/t5_encoder/t5_dense_relu_dense.py rename to src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py diff --git a/src/flux_1/models/text_encoder/t5_encoder/t5_encoder.py b/src/mflux/models/text_encoder/t5_encoder/t5_encoder.py similarity index 80% rename from src/flux_1/models/text_encoder/t5_encoder/t5_encoder.py rename to src/mflux/models/text_encoder/t5_encoder/t5_encoder.py index a010b04..b9d4e3e 100644 --- a/src/flux_1/models/text_encoder/t5_encoder/t5_encoder.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/t5_encoder/t5_feed_forward.py b/src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py similarity index 74% rename from src/flux_1/models/text_encoder/t5_encoder/t5_feed_forward.py rename to src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py index e3d7792..024e345 100644 --- a/src/flux_1/models/text_encoder/t5_encoder/t5_feed_forward.py +++ b/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): diff --git a/src/flux_1/models/text_encoder/t5_encoder/t5_layer_norm.py b/src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py similarity index 100% rename from src/flux_1/models/text_encoder/t5_encoder/t5_layer_norm.py rename to src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py diff --git a/src/flux_1/models/text_encoder/t5_encoder/t5_self_attention.py b/src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py similarity index 100% rename from src/flux_1/models/text_encoder/t5_encoder/t5_self_attention.py rename to src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py diff --git a/src/flux_1/models/transformer/__init__.py b/src/mflux/models/transformer/__init__.py similarity index 100% rename from src/flux_1/models/transformer/__init__.py rename to src/mflux/models/transformer/__init__.py diff --git a/src/flux_1/models/transformer/ada_layer_norm_continous.py b/src/mflux/models/transformer/ada_layer_norm_continous.py similarity index 95% rename from src/flux_1/models/transformer/ada_layer_norm_continous.py rename to src/mflux/models/transformer/ada_layer_norm_continous.py index eadc253..becb08d 100644 --- a/src/flux_1/models/transformer/ada_layer_norm_continous.py +++ b/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): diff --git a/src/flux_1/models/transformer/ada_layer_norm_zero.py b/src/mflux/models/transformer/ada_layer_norm_zero.py similarity index 100% rename from src/flux_1/models/transformer/ada_layer_norm_zero.py rename to src/mflux/models/transformer/ada_layer_norm_zero.py diff --git a/src/flux_1/models/transformer/ada_layer_norm_zero_single.py b/src/mflux/models/transformer/ada_layer_norm_zero_single.py similarity index 100% rename from src/flux_1/models/transformer/ada_layer_norm_zero_single.py rename to src/mflux/models/transformer/ada_layer_norm_zero_single.py diff --git a/src/flux_1/models/transformer/embed_nd.py b/src/mflux/models/transformer/embed_nd.py similarity index 100% rename from src/flux_1/models/transformer/embed_nd.py rename to src/mflux/models/transformer/embed_nd.py diff --git a/src/flux_1/models/transformer/feed_forward.py b/src/mflux/models/transformer/feed_forward.py similarity index 100% rename from src/flux_1/models/transformer/feed_forward.py rename to src/mflux/models/transformer/feed_forward.py diff --git a/src/flux_1/models/transformer/guidance_embedder.py b/src/mflux/models/transformer/guidance_embedder.py similarity index 100% rename from src/flux_1/models/transformer/guidance_embedder.py rename to src/mflux/models/transformer/guidance_embedder.py diff --git a/src/flux_1/models/transformer/joint_attention.py b/src/mflux/models/transformer/joint_attention.py similarity index 100% rename from src/flux_1/models/transformer/joint_attention.py rename to src/mflux/models/transformer/joint_attention.py diff --git a/src/flux_1/models/transformer/joint_transformer_block.py b/src/mflux/models/transformer/joint_transformer_block.py similarity index 92% rename from src/flux_1/models/transformer/joint_transformer_block.py rename to src/mflux/models/transformer/joint_transformer_block.py index 0558099..3083d53 100644 --- a/src/flux_1/models/transformer/joint_transformer_block.py +++ b/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): diff --git a/src/flux_1/models/transformer/single_block_attention.py b/src/mflux/models/transformer/single_block_attention.py similarity index 100% rename from src/flux_1/models/transformer/single_block_attention.py rename to src/mflux/models/transformer/single_block_attention.py diff --git a/src/flux_1/models/transformer/single_transformer_block.py b/src/mflux/models/transformer/single_transformer_block.py similarity index 87% rename from src/flux_1/models/transformer/single_transformer_block.py rename to src/mflux/models/transformer/single_transformer_block.py index 0c11081..4b51cf9 100644 --- a/src/flux_1/models/transformer/single_transformer_block.py +++ b/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): diff --git a/src/flux_1/models/transformer/text_embedder.py b/src/mflux/models/transformer/text_embedder.py similarity index 100% rename from src/flux_1/models/transformer/text_embedder.py rename to src/mflux/models/transformer/text_embedder.py diff --git a/src/flux_1/models/transformer/time_text_embed.py b/src/mflux/models/transformer/time_text_embed.py similarity index 82% rename from src/flux_1/models/transformer/time_text_embed.py rename to src/mflux/models/transformer/time_text_embed.py index c6f768a..7333e4c 100644 --- a/src/flux_1/models/transformer/time_text_embed.py +++ b/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): diff --git a/src/flux_1/models/transformer/timestep_embedder.py b/src/mflux/models/transformer/timestep_embedder.py similarity index 100% rename from src/flux_1/models/transformer/timestep_embedder.py rename to src/mflux/models/transformer/timestep_embedder.py diff --git a/src/flux_1/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py similarity index 87% rename from src/flux_1/models/transformer/transformer.py rename to src/mflux/models/transformer/transformer.py index 6f6b89f..ee7d095 100644 --- a/src/flux_1/models/transformer/transformer.py +++ b/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): diff --git a/src/flux_1/models/vae/__init__.py b/src/mflux/models/vae/__init__.py similarity index 100% rename from src/flux_1/models/vae/__init__.py rename to src/mflux/models/vae/__init__.py diff --git a/src/flux_1/models/vae/common/__init__.py b/src/mflux/models/vae/common/__init__.py similarity index 100% rename from src/flux_1/models/vae/common/__init__.py rename to src/mflux/models/vae/common/__init__.py diff --git a/src/flux_1/models/vae/common/attention.py b/src/mflux/models/vae/common/attention.py similarity index 96% rename from src/flux_1/models/vae/common/attention.py rename to src/mflux/models/vae/common/attention.py index 47cfb06..16e8546 100644 --- a/src/flux_1/models/vae/common/attention.py +++ b/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): diff --git a/src/flux_1/models/vae/common/resnet_block_2d.py b/src/mflux/models/vae/common/resnet_block_2d.py similarity index 98% rename from src/flux_1/models/vae/common/resnet_block_2d.py rename to src/mflux/models/vae/common/resnet_block_2d.py index 945879e..0bf73c5 100644 --- a/src/flux_1/models/vae/common/resnet_block_2d.py +++ b/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): diff --git a/src/flux_1/models/vae/common/unet_mid_block.py b/src/mflux/models/vae/common/unet_mid_block.py similarity index 84% rename from src/flux_1/models/vae/common/unet_mid_block.py rename to src/mflux/models/vae/common/unet_mid_block.py index 946c63b..c9bda8a 100644 --- a/src/flux_1/models/vae/common/unet_mid_block.py +++ b/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): diff --git a/src/flux_1/models/vae/decoder/__init__.py b/src/mflux/models/vae/decoder/__init__.py similarity index 100% rename from src/flux_1/models/vae/decoder/__init__.py rename to src/mflux/models/vae/decoder/__init__.py diff --git a/src/flux_1/models/vae/decoder/conv_in.py b/src/mflux/models/vae/decoder/conv_in.py similarity index 100% rename from src/flux_1/models/vae/decoder/conv_in.py rename to src/mflux/models/vae/decoder/conv_in.py diff --git a/src/flux_1/models/vae/decoder/conv_norm_out.py b/src/mflux/models/vae/decoder/conv_norm_out.py similarity index 93% rename from src/flux_1/models/vae/decoder/conv_norm_out.py rename to src/mflux/models/vae/decoder/conv_norm_out.py index 0420bd1..9194e93 100644 --- a/src/flux_1/models/vae/decoder/conv_norm_out.py +++ b/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): diff --git a/src/flux_1/models/vae/decoder/conv_out.py b/src/mflux/models/vae/decoder/conv_out.py similarity index 100% rename from src/flux_1/models/vae/decoder/conv_out.py rename to src/mflux/models/vae/decoder/conv_out.py diff --git a/src/flux_1/models/vae/decoder/decoder.py b/src/mflux/models/vae/decoder/decoder.py similarity index 66% rename from src/flux_1/models/vae/decoder/decoder.py rename to src/mflux/models/vae/decoder/decoder.py index 58eb408..c7c43f2 100644 --- a/src/flux_1/models/vae/decoder/decoder.py +++ b/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): diff --git a/src/flux_1/models/vae/decoder/up_block_1_or_2.py b/src/mflux/models/vae/decoder/up_block_1_or_2.py similarity index 88% rename from src/flux_1/models/vae/decoder/up_block_1_or_2.py rename to src/mflux/models/vae/decoder/up_block_1_or_2.py index 9ce2ae3..69c646b 100644 --- a/src/flux_1/models/vae/decoder/up_block_1_or_2.py +++ b/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): diff --git a/src/flux_1/models/vae/decoder/up_block_3.py b/src/mflux/models/vae/decoder/up_block_3.py similarity index 88% rename from src/flux_1/models/vae/decoder/up_block_3.py rename to src/mflux/models/vae/decoder/up_block_3.py index e25fec4..f7662bb 100644 --- a/src/flux_1/models/vae/decoder/up_block_3.py +++ b/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): diff --git a/src/flux_1/models/vae/decoder/up_block_4.py b/src/mflux/models/vae/decoder/up_block_4.py similarity index 92% rename from src/flux_1/models/vae/decoder/up_block_4.py rename to src/mflux/models/vae/decoder/up_block_4.py index b599d54..9fe9d4f 100644 --- a/src/flux_1/models/vae/decoder/up_block_4.py +++ b/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): diff --git a/src/flux_1/models/vae/decoder/up_sampler.py b/src/mflux/models/vae/decoder/up_sampler.py similarity index 100% rename from src/flux_1/models/vae/decoder/up_sampler.py rename to src/mflux/models/vae/decoder/up_sampler.py diff --git a/src/flux_1/models/vae/encoder/__init__.py b/src/mflux/models/vae/encoder/__init__.py similarity index 100% rename from src/flux_1/models/vae/encoder/__init__.py rename to src/mflux/models/vae/encoder/__init__.py diff --git a/src/flux_1/models/vae/encoder/conv_in.py b/src/mflux/models/vae/encoder/conv_in.py similarity index 100% rename from src/flux_1/models/vae/encoder/conv_in.py rename to src/mflux/models/vae/encoder/conv_in.py diff --git a/src/flux_1/models/vae/encoder/conv_norm_out.py b/src/mflux/models/vae/encoder/conv_norm_out.py similarity index 100% rename from src/flux_1/models/vae/encoder/conv_norm_out.py rename to src/mflux/models/vae/encoder/conv_norm_out.py diff --git a/src/flux_1/models/vae/encoder/conv_out.py b/src/mflux/models/vae/encoder/conv_out.py similarity index 100% rename from src/flux_1/models/vae/encoder/conv_out.py rename to src/mflux/models/vae/encoder/conv_out.py diff --git a/src/flux_1/models/vae/encoder/down_block_1.py b/src/mflux/models/vae/encoder/down_block_1.py similarity index 85% rename from src/flux_1/models/vae/encoder/down_block_1.py rename to src/mflux/models/vae/encoder/down_block_1.py index 09d3c4a..04b1b8d 100644 --- a/src/flux_1/models/vae/encoder/down_block_1.py +++ b/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): diff --git a/src/flux_1/models/vae/encoder/down_block_2.py b/src/mflux/models/vae/encoder/down_block_2.py similarity index 86% rename from src/flux_1/models/vae/encoder/down_block_2.py rename to src/mflux/models/vae/encoder/down_block_2.py index fb261c6..81542c9 100644 --- a/src/flux_1/models/vae/encoder/down_block_2.py +++ b/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): diff --git a/src/flux_1/models/vae/encoder/down_block_3.py b/src/mflux/models/vae/encoder/down_block_3.py similarity index 86% rename from src/flux_1/models/vae/encoder/down_block_3.py rename to src/mflux/models/vae/encoder/down_block_3.py index 86eddb9..6526e46 100644 --- a/src/flux_1/models/vae/encoder/down_block_3.py +++ b/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): diff --git a/src/flux_1/models/vae/encoder/down_block_4.py b/src/mflux/models/vae/encoder/down_block_4.py similarity index 89% rename from src/flux_1/models/vae/encoder/down_block_4.py rename to src/mflux/models/vae/encoder/down_block_4.py index b8eb214..fbed7fe 100644 --- a/src/flux_1/models/vae/encoder/down_block_4.py +++ b/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): diff --git a/src/flux_1/models/vae/encoder/down_sampler.py b/src/mflux/models/vae/encoder/down_sampler.py similarity index 100% rename from src/flux_1/models/vae/encoder/down_sampler.py rename to src/mflux/models/vae/encoder/down_sampler.py diff --git a/src/flux_1/models/vae/encoder/encoder.py b/src/mflux/models/vae/encoder/encoder.py similarity index 62% rename from src/flux_1/models/vae/encoder/encoder.py rename to src/mflux/models/vae/encoder/encoder.py index 14d3898..594a760 100644 --- a/src/flux_1/models/vae/encoder/encoder.py +++ b/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): diff --git a/src/flux_1/models/vae/vae.py b/src/mflux/models/vae/vae.py similarity index 85% rename from src/flux_1/models/vae/vae.py rename to src/mflux/models/vae/vae.py index 4b0402e..280afbd 100644 --- a/src/flux_1/models/vae/vae.py +++ b/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): diff --git a/src/flux_1/post_processing/__init__.py b/src/mflux/post_processing/__init__.py similarity index 100% rename from src/flux_1/post_processing/__init__.py rename to src/mflux/post_processing/__init__.py diff --git a/src/flux_1/post_processing/image.py b/src/mflux/post_processing/image.py similarity index 98% rename from src/flux_1/post_processing/image.py rename to src/mflux/post_processing/image.py index 542bfc7..5d48ffe 100644 --- a/src/flux_1/post_processing/image.py +++ b/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__) diff --git a/src/flux_1/post_processing/image_util.py b/src/mflux/post_processing/image_util.py similarity index 95% rename from src/flux_1/post_processing/image_util.py rename to src/mflux/post_processing/image_util.py index b5edd08..f6f9ea6 100644 --- a/src/flux_1/post_processing/image_util.py +++ b/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: diff --git a/src/flux_1/tokenizer/__init__.py b/src/mflux/tokenizer/__init__.py similarity index 100% rename from src/flux_1/tokenizer/__init__.py rename to src/mflux/tokenizer/__init__.py diff --git a/src/flux_1/tokenizer/clip_tokenizer.py b/src/mflux/tokenizer/clip_tokenizer.py similarity index 100% rename from src/flux_1/tokenizer/clip_tokenizer.py rename to src/mflux/tokenizer/clip_tokenizer.py diff --git a/src/flux_1/tokenizer/t5_tokenizer.py b/src/mflux/tokenizer/t5_tokenizer.py similarity index 100% rename from src/flux_1/tokenizer/t5_tokenizer.py rename to src/mflux/tokenizer/t5_tokenizer.py diff --git a/src/flux_1/tokenizer/tokenizer_handler.py b/src/mflux/tokenizer/tokenizer_handler.py similarity index 94% rename from src/flux_1/tokenizer/tokenizer_handler.py rename to src/mflux/tokenizer/tokenizer_handler.py index 5014ee7..7640e6c 100644 --- a/src/flux_1/tokenizer/tokenizer_handler.py +++ b/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: diff --git a/src/flux_1/weights/__init__.py b/src/mflux/weights/__init__.py similarity index 100% rename from src/flux_1/weights/__init__.py rename to src/mflux/weights/__init__.py diff --git a/src/flux_1/weights/lora_util.py b/src/mflux/weights/lora_util.py similarity index 98% rename from src/flux_1/weights/lora_util.py rename to src/mflux/weights/lora_util.py index afb539d..43d3fd8 100644 --- a/src/flux_1/weights/lora_util.py +++ b/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) diff --git a/src/flux_1/weights/weight_handler.py b/src/mflux/weights/weight_handler.py similarity index 98% rename from src/flux_1/weights/weight_handler.py rename to src/mflux/weights/weight_handler.py index 0314437..fc485ca 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/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: diff --git a/src/flux_1/weights/weight_util.py b/src/mflux/weights/weight_util.py similarity index 90% rename from src/flux_1/weights/weight_util.py rename to src/mflux/weights/weight_util.py index 50d8359..dd16ad3 100644 --- a/src/flux_1/weights/weight_util.py +++ b/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: