Add Qwen Image support (#269)

README.md

@@ -1591,16 +1591,16 @@ As of release [v.0.5.0](https://github.com/filipstrand/mflux/releases/tag/v.0.5.
 
 #### Training configuration
 
-To describe a training run, you need to provide a [training configuration](src/mflux/dreambooth/_example/train.json) file which specifies the details such as
+To describe a training run, you need to provide a [training configuration](src/mflux/models/flux/variants/dreambooth/_example/train.json) file which specifies the details such as
 what training data to use and various parameters. To try it out, one of the easiest ways is to start from the 
-provided [example configuration](src/mflux/dreambooth/_example/train.json) and simply use your own dataset and prompts by modifying the `examples` section of the json file.
+provided [example configuration](src/mflux/models/flux/variants/dreambooth/_example/train.json) and simply use your own dataset and prompts by modifying the `examples` section of the json file.
 
 #### Training example
 
-A complete example ([training configuration](src/mflux/dreambooth/_example/train.json) + [dataset](src/mflux/dreambooth/_example/images)) is provided in this repository. To start a training run, go to the project folder `cd mflux`, and simply run:
+A complete example ([training configuration](src/mflux/models/flux/variants/dreambooth/_example/train.json) + [dataset](src/mflux/models/flux/variants/dreambooth/_example/images)) is provided in this repository. To start a training run, go to the project folder `cd mflux`, and simply run:
 
 ```sh
-mflux-train --train-config src/mflux/dreambooth/_example/train.json
+mflux-train --train-config src/mflux/models/flux/variants/dreambooth/_example/train.json
 ```
 
 By default, this will train an adapter with images of size `512x512` with a batch size of 1 and can take up to several hours to fully complete depending on your machine.
@@ -1717,7 +1717,7 @@ To avoid this, consider some of the following strategies to reduce memory requir
 - Use a smaller batch size, for example `"batch_size": 1`
 - Make sure your Mac is not busy with other background tasks that holds memory.
 
-Applying some of these strategies, like how [train.json](src/mflux/dreambooth/_example/train.json) is set up by default,
+Applying some of these strategies, like how [train.json](src/mflux/models/flux/variants/dreambooth/_example/train.json) is set up by default,
 will allow a 32GB M1 Pro to perform a successful fine-tuning run. 
 Note, however, that reducing the trainable parameters might lead to worse performance.
  
@@ -1738,7 +1738,7 @@ The aim is to also gradually expand the scope of this feature with alternative t
 - Loss curve can be a bit misleading/hard to read, sometimes it conveys little improvement over time, but actual image samples show the real progress.
 - When plotting the loss during training, we label it as "validation loss" but it is actually only the first 10 elements of the training examples for now. Future updates should support user inputs of separate validation images.
 - Training also works with the original model as quantized!
-- [For the curious, a motivation for the loss function can be found here](src/mflux/dreambooth/optimization/_loss_derivation/dreambooth_loss.pdf).
+- [For the curious, a motivation for the loss function can be found here](src/mflux/models/flux/variants/dreambooth/optimization/_loss_derivation/dreambooth_loss.pdf).
 - Two great resources that heavily inspired this feature are: 
   - The fine-tuning script in [mlx-examples](https://github.com/ml-explore/mlx-examples/tree/main/flux#finetuning)
   - The original fine-tuning script in [Diffusers](https://huggingface.co/docs/diffusers/v0.11.0/en/training/dreambooth)

pyproject.toml

@@ -10,7 +10,7 @@ source-exclude = [
     # documentation assets were 27MB on 2025-07-05
     "**/assets/**",
     # dreambooth examples ~=5MB on 2025-07-05
-    "**/dreambooth/_example/images/**",
+    "**/models/flux/variants/dreambooth/_example/images/**",
     # loss pdf/tex does not need to be distributed
     "**/optimization/_loss_derivation/**",
 ]
@@ -26,25 +26,31 @@ maintainers = [{ name = "Filip Strand", email = "strand.filip@gmail.com" }]
 license = { file = "LICENSE" }
 requires-python = ">=3.10"
 dependencies = [
+    "accelerate>=0.31.0",
+    "filelock>=3.18.0",
     "huggingface-hub>=0.24.5,<1.0",
     "matplotlib>=3.9.2,<4.0",
     "mlx>=0.27.0,<0.28.0",
     "numpy>=2.0.1,<3.0",
     "opencv-python>=4.10.0,<5.0",
     "piexif>=1.1.3,<2.0",
+    "pillow>=10.4.0",
     "pillow>=10.4.0,<11.0; python_version<'3.13'",
     "pillow>=11.0,<12.0; python_version>='3.13'",
     "platformdirs>=4.0,<5.0",
+    "regex>=2024.11.6",
+    "requests>=2.32.4",
     "safetensors>=0.4.4,<1.0",
     "sentencepiece>=0.2.0,<1.0; python_version<'3.13'",
     # sentencepiece 0.2.1 is first release with 3.13 and 3.14 wheels: https://pypi.org/project/sentencepiece/0.2.1/
     "sentencepiece>=0.2.1,<1.0; python_version>='3.13'",
     "tokenizers>=0.20.3; python_version>='3.13'", # transformers -> tokenizers
     "toml>=0.10.2,<1.0",
+    "torch>=2.7.1",
     "torch>=2.3.1,<3.0; python_version<'3.13'",
     "torch>=2.8.0,<3.0; python_version>='3.13'",
     "tqdm>=4.66.5,<5.0",
-    "transformers>=4.44.0,<5.0",
+    "transformers>=4.55.0,<5.0",
     "twine>=6.1.0,<7.0",
 ]
 classifiers = [
@@ -84,7 +90,7 @@ mflux-save-depth = "mflux.save_depth:main"
 mflux-train = "mflux.train:main"
 mflux-upscale = "mflux.upscale:main"
 mflux-lora-library = "mflux.lora_library:main"
-mflux-completions = "mflux.completions.install:main"
+mflux-completions = "mflux.ui.cli.completions.install:main"
 
 
 [tool.ruff]

src/mflux/callbacks/callback_manager.py

@@ -12,7 +12,7 @@ class CallbackManager:
     @staticmethod
     def register_callbacks(
         args: Namespace,
-        flux,
+        model,
         enable_canny_saver: bool = False,
         enable_depth_saver: bool = False,
     ) -> MemorySaver | None:
@@ -20,7 +20,7 @@ class CallbackManager:
         CallbackManager._register_battery_saver(args)
 
         # VAE Tiling (if requested)
-        CallbackManager._register_vae_tiling(args, flux)
+        CallbackManager._register_vae_tiling(args, model)
 
         # Specialized savers (based on flags)
         if enable_canny_saver:
@@ -30,10 +30,10 @@ class CallbackManager:
             CallbackManager._register_depth_saver(args)
 
         # Stepwise handler (if requested)
-        CallbackManager._register_stepwise_handler(args, flux)
+        CallbackManager._register_stepwise_handler(args, model)
 
         # Memory saver (if requested)
-        return CallbackManager._register_memory_saver(args, flux)
+        return CallbackManager._register_memory_saver(args, model)
 
     @staticmethod
     def _register_battery_saver(args: Namespace) -> None:
@@ -41,10 +41,10 @@ class CallbackManager:
         CallbackRegistry.register_before_loop(battery_saver)
 
     @staticmethod
-    def _register_vae_tiling(args: Namespace, flux) -> None:
+    def _register_vae_tiling(args: Namespace, model) -> None:
         if args.vae_tiling:
-            flux.vae.decoder.enable_tiling = True
-            flux.vae.decoder.split_direction = args.vae_tiling_split
+            model.vae.decoder.enable_tiling = True
+            model.vae.decoder.split_direction = args.vae_tiling_split
 
     @staticmethod
     def _register_canny_saver(args: Namespace) -> None:
@@ -59,18 +59,18 @@ class CallbackManager:
             CallbackRegistry.register_before_loop(depth_image_saver)
 
     @staticmethod
-    def _register_stepwise_handler(args: Namespace, flux) -> None:
+    def _register_stepwise_handler(args: Namespace, model) -> None:
         if args.stepwise_image_output_dir:
-            handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+            handler = StepwiseHandler(model=model, output_dir=args.stepwise_image_output_dir)
             CallbackRegistry.register_before_loop(handler)
             CallbackRegistry.register_in_loop(handler)
             CallbackRegistry.register_interrupt(handler)
 
     @staticmethod
-    def _register_memory_saver(args: Namespace, flux) -> MemorySaver | None:
+    def _register_memory_saver(args: Namespace, model) -> MemorySaver | None:
         memory_saver = None
         if args.low_ram:
-            memory_saver = MemorySaver(flux=flux, keep_transformer=len(args.seed) > 1)
+            memory_saver = MemorySaver(model=model, keep_transformer=len(args.seed) > 1)
             CallbackRegistry.register_before_loop(memory_saver)
             CallbackRegistry.register_in_loop(memory_saver)
             CallbackRegistry.register_after_loop(memory_saver)

src/mflux/callbacks/instances/memory_saver.py

@@ -14,8 +14,8 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback):
     components at strategic points in the execution cycle.
     """
 
-    def __init__(self, flux, keep_transformer: bool = True, cache_limit_bytes: int = 1000**3):
-        self.flux = flux
+    def __init__(self, model, keep_transformer: bool = True, cache_limit_bytes: int = 1000**3):
+        self.model = model
         self.keep_transformer = keep_transformer
         self.peak_memory: int = 0
         mx.set_cache_limit(cache_limit_bytes)
@@ -32,7 +32,7 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback):
         depth_image: PIL.Image.Image | None = None,
     ) -> None:
         self.peak_memory = mx.get_peak_memory()
-        self._delete_encoders()
+        self._delete_text_encoders()
 
     def call_in_loop(
         self,
@@ -56,17 +56,21 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback):
         if not self.keep_transformer:
             self._delete_transformer()
 
-    def _delete_encoders(self) -> None:
+    def _delete_text_encoders(self) -> None:
         # repeated image generation only works with the same prompt (cache)
-        self.flux.clip_text_encoder = None
-        self.flux.t5_text_encoder = None
+        if hasattr(self.model, "clip_image_encoder"):
+            self.model.clip_image_encoder = None
+        if hasattr(self.model, "t5_image_encoder"):
+            self.model.t5_image_encoder = None
+        if hasattr(self.model, "text_encoder"):
+            self.model.text_encoder = None
         gc.collect()
         mx.clear_cache()
 
     def _delete_transformer(self) -> None:
-        self.flux.transformer = None
-        if hasattr(self.flux, "transformer_controlnet"):
-            self.flux.transformer_controlnet = None
+        self.model.transformer = None
+        if hasattr(self.model, "transformer_controlnet"):
+            self.model.transformer_controlnet = None
         gc.collect()
         mx.clear_cache()
 

src/mflux/callbacks/instances/stepwise_handler.py

@@ -13,10 +13,10 @@ from mflux.post_processing.image_util import ImageUtil
 class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback):
     def __init__(
         self,
-        flux,
+        model,
         output_dir: str,
     ):
-        self.flux = flux
+        self.model = model
         self.output_dir = Path(output_dir)
         self.step_wise_images = []
 
@@ -80,16 +80,16 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback):
         time_steps: tqdm,
     ) -> None:
         unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
-        stepwise_decoded = self.flux.vae.decode(unpack_latents)
+        stepwise_decoded = self.model.vae.decode(unpack_latents)
         generation_time = time_steps.format_dict["elapsed"] if time_steps is not None else 0
         stepwise_img = ImageUtil.to_image(
             decoded_latents=stepwise_decoded,
             config=config,
             seed=seed,
             prompt=prompt,
-            quantization=self.flux.bits,
-            lora_paths=self.flux.lora_paths,
-            lora_scales=self.flux.lora_scales,
+            quantization=self.model.bits,
+            lora_paths=self.model.lora_paths,
+            lora_scales=self.model.lora_scales,
             generation_time=generation_time,
         )
         stepwise_img.save(

src/mflux/community/concept_attention/__init__.py

@@ -1 +0,0 @@
-# Concept Attention transformer models

src/mflux/concept.py

@@ -1,8 +1,8 @@
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.community.concept_attention.flux_concept import Flux1Concept
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.flux.variants.concept_attention.flux_concept import Flux1Concept
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -34,7 +34,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux)
 
     try:
         for seed in args.seed:

src/mflux/concept_from_image.py

@@ -1,8 +1,8 @@
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.community.concept_attention.flux_concept_from_image import Flux1ConceptFromImage
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.flux.variants.concept_attention.flux_concept_from_image import Flux1ConceptFromImage
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -34,7 +34,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux)
 
     try:
         for seed in args.seed:

src/mflux/config/model_config.py

@@ -84,6 +84,11 @@ class ModelConfig:
     def krea_dev() -> "ModelConfig":
         return AVAILABLE_MODELS["krea-dev"]
 
+    @staticmethod
+    @lru_cache
+    def qwen_image() -> "ModelConfig":
+        return AVAILABLE_MODELS["qwen-image"]
+
     def x_embedder_input_dim(self) -> int:
         if "Fill" in self.model_name:
             return 384
@@ -281,4 +286,16 @@ AVAILABLE_MODELS = {
         requires_sigma_shift=True,
         priority=10,
     ),
+    "qwen-image": ModelConfig(
+        aliases=["qwen-image", "qwen"],
+        model_name="Qwen/Qwen-Image",
+        base_model=None,
+        controlnet_model=None,
+        custom_transformer_model=None,
+        num_train_steps=None,
+        max_sequence_length=None,
+        supports_guidance=None,
+        requires_sigma_shift=None,
+        priority=11,
+    ),
 }

src/mflux/generate.py

@@ -2,10 +2,11 @@ from mflux.callbacks.callback_manager import CallbackManager
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
-from mflux.flux.flux import Flux1
+from mflux.models.flux.variants.txt2img.flux import Flux1
+from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
-from mflux.ui.prompt_utils import get_effective_prompt
+from mflux.ui.prompt_utils import get_effective_negative_prompt, get_effective_prompt
 
 
 def main():
@@ -24,7 +25,8 @@ def main():
         args.guidance = ui_defaults.GUIDANCE_SCALE
 
     # 1. Load the model
-    flux = Flux1(
+    model_class = QwenImage if "qwen" in args.model.lower() else Flux1
+    model = model_class(
         model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model),
         quantize=args.quantize,
         local_path=args.path,
@@ -33,14 +35,15 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=model)
 
     try:
         for seed in args.seed:
             # 3. Generate an image for each seed value
-            image = flux.generate_image(
+            image = model.generate_image(
                 seed=seed,
                 prompt=get_effective_prompt(args),
+                negative_prompt=get_effective_negative_prompt(args),
                 config=Config(
                     num_inference_steps=args.steps,
                     height=args.height,

src/mflux/generate_controlnet.py

@@ -1,8 +1,8 @@
 from mflux.callbacks.callback_manager import CallbackManager
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
-from mflux.controlnet.flux_controlnet import Flux1Controlnet
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.flux.variants.controlnet.flux_controlnet import Flux1Controlnet
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -33,7 +33,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux, enable_canny_saver=True)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux, enable_canny_saver=True)
 
     try:
         for seed in args.seed:

src/mflux/generate_depth.py

@@ -1,7 +1,7 @@
 from mflux.callbacks.callback_manager import CallbackManager
 from mflux.config.config import Config
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
-from mflux.flux_tools.depth.flux_depth import Flux1Depth
+from mflux.models.flux.variants.depth.flux_depth import Flux1Depth
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -31,7 +31,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux, enable_depth_saver=True)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux, enable_depth_saver=True)
 
     try:
         for seed in args.seed:

src/mflux/generate_fill.py

@@ -1,7 +1,7 @@
 from mflux.callbacks.callback_manager import CallbackManager
 from mflux.config.config import Config
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
-from mflux.flux_tools.fill.flux_fill import Flux1Fill
+from mflux.models.flux.variants.fill.flux_fill import Flux1Fill
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -31,7 +31,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux)
 
     try:
         for seed in args.seed:

src/mflux/generate_in_context_catvton.py

@@ -1,10 +1,10 @@
 from pathlib import Path
 
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.community.in_context.flux_in_context_fill import Flux1InContextFill
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -44,7 +44,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux)
 
     try:
         for seed in args.seed:

src/mflux/generate_in_context_dev.py

@@ -1,11 +1,11 @@
 from pathlib import Path
 
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.community.in_context.flux_in_context_dev import Flux1InContextDev
-from mflux.community.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.flux.variants.in_context.flux_in_context_dev import Flux1InContextDev
+from mflux.models.flux.variants.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -42,7 +42,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux)
 
     try:
         for seed in args.seed:

src/mflux/generate_in_context_edit.py

@@ -3,11 +3,11 @@ from pathlib import Path
 from PIL import Image
 
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.community.in_context.flux_in_context_fill import Flux1InContextFill
-from mflux.community.in_context.utils.in_context_loras import prepare_ic_edit_loras
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill
+from mflux.models.flux.variants.in_context.utils.in_context_loras import prepare_ic_edit_loras
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -46,7 +46,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux)
 
     try:
         for seed in args.seed:

src/mflux/generate_kontext.py

@@ -3,7 +3,7 @@ from pathlib import Path
 from mflux.callbacks.callback_manager import CallbackManager
 from mflux.config.config import Config
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
-from mflux.kontext.flux_kontext import Flux1Kontext
+from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -33,7 +33,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux)
 
     try:
         for seed in args.seed:

src/mflux/generate_redux.py

@@ -4,7 +4,7 @@ from mflux.callbacks.callback_manager import CallbackManager
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
-from mflux.flux_tools.redux.flux_redux import Flux1Redux
+from mflux.models.flux.variants.redux.flux_redux import Flux1Redux
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt
@@ -41,7 +41,7 @@ def main():
     )
 
     # 2. Register callbacks
-    memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
+    memory_saver = CallbackManager.register_callbacks(args=args, model=flux)
 
     try:
         for seed in args.seed:

src/mflux/latent_creator/latent_creator.py

@@ -1,8 +1,8 @@
 from pathlib import Path
 
 import mlx.core as mx
+from mlx import nn
 
-from mflux.models.vae.vae import VAE
 from mflux.post_processing.array_util import ArrayUtil
 from mflux.post_processing.image_util import ImageUtil
 
@@ -10,7 +10,7 @@ from mflux.post_processing.image_util import ImageUtil
 class Img2Img:
     def __init__(
         self,
-        vae: VAE,
+        vae: nn.Module,
         sigmas: mx.array,
         init_time_step: int,
         image_path: str | Path | None,
@@ -71,7 +71,7 @@ class LatentCreator:
             )
 
     @staticmethod
-    def encode_image(vae: VAE, image_path: str | Path, height: int, width: int):
+    def encode_image(vae: nn.Module, image_path: str | Path, height: int, width: int):
         scaled_user_image = ImageUtil.scale_to_dimensions(
             image=ImageUtil.load_image(image_path).convert("RGB"),
             target_width=width,

src/mflux/lora_library.py

@@ -7,7 +7,7 @@ import sys
 from collections import defaultdict
 from pathlib import Path
 
-from mflux.weights.lora_library import _discover_lora_files
+from mflux.utils.lora_library import _discover_lora_files
 
 
 def list_loras(paths: list[str] | None = None) -> int:

src/mflux/models/common/lora/download/lora_huggingface_downloader.py

@@ -2,28 +2,30 @@ import shutil
 from pathlib import Path
 
 from mflux.ui.defaults import MFLUX_LORA_CACHE_DIR
-from mflux.weights.download import snapshot_download
+from mflux.utils.download import snapshot_download
 
 
-class WeightHandlerLoRAHuggingFace:
+class LoRAHuggingFaceDownloader:
+
     @staticmethod
     def download_loras(
         lora_names: list[str] | None = None,
         repo_id: str | None = None,
         cache_dir: Path | str | None = None,
+        model_name: str = "LoRA",
     ) -> list[str]:
-        if repo_id is None:
+        if not lora_names or not repo_id:
             return []
 
         lora_paths = []
-        if lora_names:
-            for lora_name in lora_names:
-                lora_path = WeightHandlerLoRAHuggingFace.download_lora(
-                    repo_id=repo_id,
-                    lora_name=lora_name,
-                    cache_dir=cache_dir,
-                )
-                lora_paths.append(lora_path)
+        for lora_name in lora_names:
+            lora_path = LoRAHuggingFaceDownloader.download_lora(
+                repo_id=repo_id,
+                lora_name=lora_name,
+                cache_dir=cache_dir,
+                model_name=model_name,
+            )
+            lora_paths.append(lora_path)
 
         return lora_paths
 
@@ -32,6 +34,7 @@ class WeightHandlerLoRAHuggingFace:
         repo_id: str,
         lora_name: str,
         cache_dir: Path | str | None = None,
+        model_name: str = "LoRA",  # For logging purposes
     ) -> str:
         # Ensure cache_dir is a Path object
         if cache_dir is None:
@@ -41,25 +44,25 @@ class WeightHandlerLoRAHuggingFace:
 
         cache_path.mkdir(parents=True, exist_ok=True)
 
-        # Check if the file already exists in the cache
+        # Check if already cached
         cached_file_path = cache_path / lora_name
         if cached_file_path.exists() and cached_file_path.is_file():
             try:
                 # Verify the file is actually readable (catches broken symlinks)
                 with open(cached_file_path, "rb") as f:
                     f.read(1)  # Try to read just 1 byte to verify it works
-                print(f"Using cached LoRA: {cached_file_path}")
+                print(f"Using cached {model_name} LoRA: {cached_file_path}")
                 return str(cached_file_path)
             except (OSError, IOError):
                 # File exists but is not readable (broken symlink, permissions, etc.)
-                print(f"Cached LoRA file is corrupted or inaccessible, re-downloading: {cached_file_path}")
+                print(f"Cached {model_name} LoRA file is corrupted or inaccessible, re-downloading: {cached_file_path}")
                 try:
                     cached_file_path.unlink()  # Remove the broken file/symlink
                 except OSError:
                     pass  # Ignore if we can't remove it
 
         # Download the LoRA from Hugging Face
-        print(f"Downloading LoRA '{lora_name}' from {repo_id}...")
+        print(f"Downloading {model_name} LoRA '{lora_name}' from {repo_id}...")
         download_path = Path(
             snapshot_download(
                 repo_id=repo_id,
@@ -69,18 +72,30 @@ class WeightHandlerLoRAHuggingFace:
         )
 
         # Find the downloaded file
-        for file in download_path.glob(f"**/*{lora_name}*"):
+        print(f"🔍 Searching for downloaded files in: {download_path}")
+        found_files = list(download_path.glob(f"**/*{lora_name}*"))
+        print(f"📁 Found files matching pattern: {found_files}")
+
+        for file in found_files:
+            print(f"📄 Checking file: {file} (suffix: {file.suffix}, size: {file.stat().st_size} bytes)")
             if file.is_file() and file.suffix in [".safetensors", ".bin"]:
-                # Copy or link the file to the cache directory with the expected name
-                target_path = cache_path / lora_name
+                # Ensure the target path has the correct extension
+                if not lora_name.endswith(file.suffix):
+                    target_name = f"{lora_name}{file.suffix}"
+                else:
+                    target_name = lora_name
+
+                target_path = cache_path / target_name
                 if not target_path.exists():
                     # Create a symlink or copy the file
                     try:
                         target_path.symlink_to(file)
+                        print(f"🔗 Created symlink: {target_path} -> {file}")
                     except (OSError, AttributeError):
                         shutil.copy2(file, target_path)
+                        print(f"📋 Copied file: {file} -> {target_path}")
 
-                print(f"LoRA downloaded to: {target_path}")
+                print(f"{model_name} LoRA downloaded to: {target_path}")
                 return str(target_path)
 
-        raise FileNotFoundError(f"Could not find LoRA file '{lora_name}' in the downloaded repository.")
+        raise FileNotFoundError(f"Could not find {model_name} LoRA file '{lora_name}' in the downloaded repository.")

src/mflux/models/common/lora/layer/fused_linear_lora_layer.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from mflux.dreambooth.lora_layers.linear_lora_layer import LoRALinear
+from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
 
 
 class FusedLoRALinear(nn.Module):

src/mflux/models/common/lora/mapping/lora_loader.py

@@ -0,0 +1,248 @@
+from pathlib import Path
+from typing import Dict, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mflux.models.common.lora.layer.fused_linear_lora_layer import FusedLoRALinear
+from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
+from mflux.models.common.lora.mapping.lora_mapping import LoRATarget
+
+
+class LoRALoader:
+
+    @staticmethod
+    def load_and_apply_lora(
+        lora_mapping: list[LoRATarget],
+        transformer: nn.Module,
+        lora_files: list[str],
+        lora_scales: list[float] | None = None,
+    ) -> None:
+        if not lora_files:
+            return
+
+        # Validate scales - handle both None and empty list cases
+        if lora_scales is None or len(lora_scales) == 0:
+            lora_scales = [1.0] * len(lora_files)
+        elif len(lora_scales) != len(lora_files):
+            raise ValueError(
+                f"Number of LoRA scales ({len(lora_scales)}) must match number of LoRA files ({len(lora_files)})"
+            )
+
+        print(f"📦 Loading {len(lora_files)} LoRA file(s)...")
+
+        for lora_file, scale in zip(lora_files, lora_scales):
+            LoRALoader._apply_single_lora(transformer, lora_file, scale, lora_mapping)
+
+        print("✅ All LoRA weights applied successfully")
+
+    @staticmethod
+    def _apply_single_lora(
+        transformer: nn.Module,
+        lora_file: str,
+        scale: float,
+        lora_mapping: list[LoRATarget]
+    ) -> None:
+        # Load the LoRA weights
+        if not Path(lora_file).exists():
+            print(f"❌ LoRA file not found: {lora_file}")
+            return
+
+        print(f"🔧 Applying LoRA: {Path(lora_file).name} (scale={scale})")
+
+        try:
+            weights = dict(mx.load(lora_file, return_metadata=True)[0].items())
+            mx.eval(weights)
+        except (FileNotFoundError, ValueError, RuntimeError) as e:
+            print(f"❌ Failed to load LoRA file: {e}")
+            return
+
+        print(f"🔍 DEBUG: Found {len(weights)} LoRA weights")
+
+        # Apply LoRA using the provided mapping
+        flat_mapping = LoRALoader._get_flat_mapping(lora_mapping)
+        applied_count = LoRALoader._apply_lora_with_mapping(transformer, weights, scale, flat_mapping)
+
+        print(f"   ✅ Applied to {applied_count} layers")
+
+    @staticmethod
+    def _apply_lora_with_mapping(
+        transformer: nn.Module,
+        weights: dict,
+        scale: float,
+        lora_mappings: Dict[str, Tuple[str, str, bool]]
+    ) -> int:
+        applied_count = 0
+        lora_data_by_target = {}
+
+        # Group LoRA weights by their target layers
+        for weight_key, weight_value in weights.items():
+            found_mapping = None
+            block_idx = None
+
+            # Pattern matching logic
+            for pattern, mapping_info in lora_mappings.items():
+                if "{block}" in pattern:
+                    # Extract block number from the weight key
+                    parts = weight_key.split(".")
+                    for i, part in enumerate(parts):
+                        if part.isdigit():
+                            try:
+                                test_block_idx = int(part)
+                                concrete_pattern = pattern.format(block=test_block_idx)
+                                if weight_key == concrete_pattern:
+                                    found_mapping = mapping_info
+                                    block_idx = test_block_idx
+                                    break
+                            except (ValueError, KeyError):
+                                continue
+                    if found_mapping:
+                        break
+                else:
+                    if weight_key == pattern:
+                        found_mapping = mapping_info
+                        break
+
+            if found_mapping is None:
+                continue
+
+            target_path, matrix_name, transpose = found_mapping
+
+            # Handle block substitution in target path
+            if block_idx is not None and "{block}" in target_path:
+                target_path = target_path.format(block=block_idx)
+
+            if target_path not in lora_data_by_target:
+                lora_data_by_target[target_path] = {}
+
+            lora_data_by_target[target_path][matrix_name] = (weight_value, transpose)
+
+        print(f"🔍 DEBUG: Found {len(lora_data_by_target)} LoRA target layers")
+
+        # Apply LoRA to each target
+        for target_path, lora_data in lora_data_by_target.items():
+            if LoRALoader._apply_lora_matrices_to_target(transformer, target_path, lora_data, scale):
+                applied_count += 1
+
+        return applied_count
+
+    @staticmethod
+    def _apply_lora_matrices_to_target(
+        transformer: nn.Module,
+        target_path: str,
+        lora_data: dict,
+        scale: float
+    ) -> bool:
+        # Navigate to the target layer
+        current_module = transformer
+        path_parts = target_path.split(".")
+
+        try:
+            for part in path_parts:
+                if part.isdigit():
+                    current_module = current_module[int(part)]
+                else:
+                    current_module = getattr(current_module, part)
+        except (AttributeError, IndexError, KeyError):
+            print(f"❌ Could not find target path: {target_path}")
+            return False
+
+        # Check if we have the required matrices
+        if "lora_A" not in lora_data or "lora_B" not in lora_data:
+            print(f"❌ Missing required LoRA matrices for {target_path}")
+            return False
+
+        lora_A, transpose_A = lora_data["lora_A"]
+        lora_B, transpose_B = lora_data["lora_B"]
+
+        # Handle transposition
+        if transpose_A:
+            lora_A = lora_A.T
+        if transpose_B:
+            lora_B = lora_B.T
+
+        # Handle alpha scaling
+        alpha_scale = 1.0
+        if "alpha" in lora_data:
+            alpha_value, _ = lora_data["alpha"]
+            rank = lora_A.shape[1]
+            alpha_scale = float(alpha_value) / rank
+
+        # Calculate final scale - only use user scale, matching Diffusers approach
+        effective_scale = scale
+
+        # Create new LoRA layer
+        if hasattr(current_module, 'weight'):
+            # Create LoRA layer
+            lora_layer = LoRALinear.from_linear(
+                current_module,
+                r=lora_A.shape[1],
+                scale=effective_scale
+            )
+
+            # Set the LoRA matrices - use the correct dimensions from the LoRA file
+            lora_layer.lora_A = lora_A
+            lora_layer.lora_B = lora_B
+
+            # Apply alpha scaling to the matrices if present
+            if "alpha" in lora_data:
+                lora_layer.lora_B = lora_layer.lora_B * alpha_scale
+
+            # Handle fusion: if the current module is already a LoRA layer, fuse them
+            if isinstance(current_module, LoRALinear):
+                print(f"   🔀 Fusing with existing LoRA at {target_path}")
+                # Create fused layer with the existing LoRA and the new one
+                fused_layer = FusedLoRALinear(
+                    base_linear=current_module.linear,
+                    loras=[current_module, lora_layer]
+                )
+                replacement_layer = fused_layer
+            elif isinstance(current_module, FusedLoRALinear):
+                print(f"   🔀 Adding to existing fusion at {target_path}")
+                # Add to existing fusion
+                fused_layer = FusedLoRALinear(
+                    base_linear=current_module.base_linear,
+                    loras=current_module.loras + [lora_layer]
+                )
+                replacement_layer = fused_layer
+            else:
+                # First LoRA on this layer
+                replacement_layer = lora_layer
+
+            # Replace the layer in the parent module
+            parent_module = transformer
+            for part in path_parts[:-1]:
+                if part.isdigit():
+                    parent_module = parent_module[int(part)]
+                else:
+                    parent_module = getattr(parent_module, part)
+
+            final_attr = path_parts[-1]
+            if final_attr.isdigit():
+                parent_module[int(final_attr)] = replacement_layer
+            else:
+                setattr(parent_module, final_attr, replacement_layer)
+
+            return True
+        else:
+            print(f"❌ Target layer {target_path} is not a linear layer")
+            return False
+
+    @staticmethod
+    def _get_flat_mapping(targets: list[LoRATarget]) -> Dict[str, Tuple[str, str, bool]]:
+        flat_mapping = {}
+
+        for target in targets:
+            # Add up weight patterns (lora_B, transposed)
+            for pattern in target.possible_up_patterns:
+                flat_mapping[pattern] = (target.model_path, "lora_B", True)
+
+            # Add down weight patterns (lora_A, transposed)
+            for pattern in target.possible_down_patterns:
+                flat_mapping[pattern] = (target.model_path, "lora_A", True)
+
+            # Add alpha patterns (no transpose)
+            for pattern in target.possible_alpha_patterns:
+                flat_mapping[pattern] = (target.model_path, "alpha", False)
+
+        return flat_mapping

src/mflux/models/common/lora/mapping/lora_mapping.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from typing import List, Protocol
+
+
+@dataclass
+class LoRATarget:
+    model_path: str
+    possible_up_patterns: List[str]
+    possible_down_patterns: List[str]
+    possible_alpha_patterns: List[str]
+
+
+class LoRAMapping(Protocol):
+
+    @staticmethod
+    def get_mapping() -> List[LoRATarget]:
+        return

src/mflux/models/depth_pro/depth_pro.py

@@ -7,8 +7,8 @@ import numpy as np
 from PIL import Image
 
 from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer
-from mflux.models.depth_pro.depth_pro_model import DepthProModel
-from mflux.models.depth_pro.depth_pro_util import DepthProUtil
+from mflux.models.depth_pro.model.depth_pro_model import DepthProModel
+from mflux.models.depth_pro.model.depth_pro_util import DepthProUtil
 from mflux.post_processing.image_util import ImageUtil
 
 

src/mflux/models/depth_pro/depth_pro_initializer.py

@@ -1,7 +1,7 @@
 import mlx.nn as nn
 
-from mflux.models.depth_pro.depth_pro_model import DepthProModel
-from mflux.models.depth_pro.weight_handler_depth_pro import WeightHandlerDepthPro
+from mflux.models.depth_pro.model.depth_pro_model import DepthProModel
+from mflux.models.depth_pro.weights.weight_handler_depth_pro import WeightHandlerDepthPro
 
 
 class DepthProInitializer:

src/mflux/models/depth_pro/model/depth_pro_encoder.py

@@ -3,9 +3,9 @@ import math
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.conv_utils import ConvUtils
-from mflux.models.depth_pro.dino_v2.dino_vision_transformer import DinoVisionTransformer
-from mflux.models.depth_pro.upsample_block import UpSampleBlock
+from mflux.models.depth_pro.model.conv_utils import ConvUtils
+from mflux.models.depth_pro.model.dino_v2.dino_vision_transformer import DinoVisionTransformer
+from mflux.models.depth_pro.model.upsample_block import UpSampleBlock
 
 
 class DepthProEncoder(nn.Module):

src/mflux/models/depth_pro/model/depth_pro_model.py

@@ -1,9 +1,9 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.depth_pro_encoder import DepthProEncoder
-from mflux.models.depth_pro.fov_head import FOVHead
-from mflux.models.depth_pro.multires_conv_decoder import MultiresConvDecoder
+from mflux.models.depth_pro.model.depth_pro_encoder import DepthProEncoder
+from mflux.models.depth_pro.model.fov_head import FOVHead
+from mflux.models.depth_pro.model.multires_conv_decoder import MultiresConvDecoder
 
 
 class DepthProModel(nn.Module):

src/mflux/models/depth_pro/model/dino_v2/dino_vision_transformer.py

@@ -1,8 +1,8 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.dino_v2.patch_embed import PatchEmbed
-from mflux.models.depth_pro.dino_v2.transformer_block import TransformerBlock
+from mflux.models.depth_pro.model.dino_v2.patch_embed import PatchEmbed
+from mflux.models.depth_pro.model.dino_v2.transformer_block import TransformerBlock
 
 
 class DinoVisionTransformer(nn.Module):

src/mflux/models/depth_pro/model/dino_v2/transformer_block.py

@@ -1,9 +1,9 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.dino_v2.attention import Attention
-from mflux.models.depth_pro.dino_v2.layer_scale import LayerScale
-from mflux.models.depth_pro.dino_v2.mlp import MLP
+from mflux.models.depth_pro.model.dino_v2.attention import Attention
+from mflux.models.depth_pro.model.dino_v2.layer_scale import LayerScale
+from mflux.models.depth_pro.model.dino_v2.mlp import MLP
 
 
 class TransformerBlock(nn.Module):

src/mflux/models/depth_pro/model/feature_fusion_block_2d.py

@@ -1,8 +1,8 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.conv_utils import ConvUtils
-from mflux.models.depth_pro.residual_block import ResidualBlock
+from mflux.models.depth_pro.model.conv_utils import ConvUtils
+from mflux.models.depth_pro.model.residual_block import ResidualBlock
 
 
 class FeatureFusionBlock2d(nn.Module):

src/mflux/models/depth_pro/model/fov_head.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.conv_utils import ConvUtils
+from mflux.models.depth_pro.model.conv_utils import ConvUtils
 
 
 class FOVHead(nn.Module):

src/mflux/models/depth_pro/model/multires_conv_decoder.py

@@ -1,8 +1,8 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.conv_utils import ConvUtils
-from mflux.models.depth_pro.feature_fusion_block_2d import FeatureFusionBlock2d
+from mflux.models.depth_pro.model.conv_utils import ConvUtils
+from mflux.models.depth_pro.model.feature_fusion_block_2d import FeatureFusionBlock2d
 
 
 class MultiresConvDecoder(nn.Module):

src/mflux/models/depth_pro/model/residual_block.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.conv_utils import ConvUtils
+from mflux.models.depth_pro.model.conv_utils import ConvUtils
 
 
 class ResidualBlock(nn.Module):

src/mflux/models/depth_pro/model/upsample_block.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from mflux.models.depth_pro.conv_utils import ConvUtils
+from mflux.models.depth_pro.model.conv_utils import ConvUtils
 
 
 class UpSampleBlock(nn.Module):

src/mflux/models/depth_pro/weights/weight_handler_depth_pro.py

@@ -6,9 +6,9 @@ import mlx.core as mx
 import torch
 from mlx.utils import tree_unflatten
 
+from mflux.models.flux.weights.weight_handler import MetaData
+from mflux.models.flux.weights.weight_util import WeightUtil
 from mflux.ui.defaults import MFLUX_CACHE_DIR
-from mflux.weights.weight_handler import MetaData
-from mflux.weights.weight_util import WeightUtil
 
 logger = logging.getLogger(__name__)
 

src/mflux/models/flux/flux_initializer.py

@@ -1,21 +1,21 @@
 from mflux.config.model_config import ModelConfig
-from mflux.controlnet.transformer_controlnet import TransformerControlnet
-from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
-from mflux.flux_tools.redux.weight_handler_redux import WeightHandlerRedux
 from mflux.models.depth_pro.depth_pro import DepthPro
-from mflux.models.redux_encoder.redux_encoder import ReduxEncoder
-from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer
-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.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
-from mflux.weights.weight_handler_lora import WeightHandlerLoRA
-from mflux.weights.weight_handler_lora_huggingface import WeightHandlerLoRAHuggingFace
-from mflux.weights.weight_util import WeightUtil
+from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder
+from mflux.models.flux.model.flux_transformer.transformer import Transformer
+from mflux.models.flux.model.flux_vae.vae import VAE
+from mflux.models.flux.model.redux_encoder.redux_encoder import ReduxEncoder
+from mflux.models.flux.model.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer
+from mflux.models.flux.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.models.flux.tokenizer.t5_tokenizer import TokenizerT5
+from mflux.models.flux.tokenizer.tokenizer_handler import TokenizerHandler
+from mflux.models.flux.variants.controlnet.transformer_controlnet import TransformerControlnet
+from mflux.models.flux.variants.controlnet.weight_handler_controlnet import WeightHandlerControlnet
+from mflux.models.flux.variants.redux.weight_handler_redux import WeightHandlerRedux
+from mflux.models.common.lora.download.lora_huggingface_downloader import LoRAHuggingFaceDownloader
+from mflux.models.flux.weights.weight_handler import WeightHandler
+from mflux.models.flux.weights.weight_handler_lora import WeightHandlerLoRA
+from mflux.models.flux.weights.weight_util import WeightUtil
 
 
 class FluxInitializer:
@@ -81,9 +81,10 @@ class FluxInitializer:
         )
 
         # 5. Set LoRA weights
-        hf_lora_paths = WeightHandlerLoRAHuggingFace.download_loras(
+        hf_lora_paths = LoRAHuggingFaceDownloader.download_loras(
             lora_names=lora_names,
             repo_id=lora_repo_id,
+            model_name="Flux",
         )
         flux_model.lora_paths = lora_paths + hf_lora_paths
         flux_model.lora_scales = (lora_scales or []) + [1.0] * len(hf_lora_paths)
@@ -206,9 +207,9 @@ class FluxInitializer:
         lora_repo_id: str | None = None,
     ):
         # Import here to avoid circular dependency
-        from mflux.community.concept_attention.transformer_concept import TransformerConcept
+        from mflux.models.flux.variants.concept_attention.transformer_concept import TransformerConcept
 
-        # 1. Load weights first to get transformer dimensions
+        # 1. Load weights first to get flux_transformer dimensions
         weights = WeightHandler.load_regular_weights(
             repo_id=model_config.model_name,
             local_path=local_path,
@@ -221,7 +222,7 @@ class FluxInitializer:
             num_single_transformer_blocks=weights.num_single_transformer_blocks(),
         )
 
-        # 3. Use the improved FluxInitializer with custom transformer
+        # 3. Use the improved FluxInitializer with custom flux_transformer
         FluxInitializer.init(
             flux_model=flux_model,
             model_config=model_config,

src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_embeddings.py

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

src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_encoder.py

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

src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_encoder_layer.py

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

src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_text_model.py

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

src/mflux/models/flux/model/flux_text_encoder/clip_encoder/encoder_clip.py

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

src/mflux/models/flux/model/flux_text_encoder/prompt_encoder.py

@@ -1,9 +1,9 @@
 import mlx.core as mx
 
-from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
-from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
-from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
-from mflux.tokenizer.t5_tokenizer import TokenizerT5
+from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder
+from mflux.models.flux.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.models.flux.tokenizer.t5_tokenizer import TokenizerT5
 
 
 class PromptEncoder:

src/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_attention.py

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

src/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_block.py

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

src/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_encoder.py

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

src/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_feed_forward.py

@@ -1,10 +1,8 @@
 import mlx.core as mx
 from mlx import nn
 
-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
+from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_dense_relu_dense import T5DenseReluDense
+from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
 
 
 class T5FeedForward(nn.Module):

src/mflux/models/flux/model/flux_transformer/common/attention_utils.py

@@ -0,0 +1,98 @@
+import mlx.core as mx
+from mlx import nn
+from mlx.core.fast import scaled_dot_product_attention
+
+
+class AttentionUtils:
+    @staticmethod
+    def process_qkv(
+        hidden_states: mx.array,
+        to_q: nn.Linear,
+        to_k: nn.Linear,
+        to_v: nn.Linear,
+        norm_q: nn.RMSNorm,
+        norm_k: nn.RMSNorm,
+        num_heads: int,
+        head_dim: int,
+    ):
+        batch_size = hidden_states.shape[0]
+        seq_len = hidden_states.shape[1]
+
+        query = to_q(hidden_states)
+        key = to_k(hidden_states)
+        value = to_v(hidden_states)
+
+        # Reshape [B, S, H*D] -> [B, H, S, D]
+        query = mx.transpose(mx.reshape(query, (batch_size, seq_len, num_heads, head_dim)), (0, 2, 1, 3))
+        key = mx.transpose(mx.reshape(key, (batch_size, seq_len, num_heads, head_dim)), (0, 2, 1, 3))
+        value = mx.transpose(mx.reshape(value, (batch_size, seq_len, num_heads, head_dim)), (0, 2, 1, 3))
+
+        # Apply normalization in float32 for stability, then cast back (like working version)
+        q_dtype = query.dtype
+        k_dtype = key.dtype
+        query = norm_q(query.astype(mx.float32)).astype(q_dtype)
+        key = norm_k(key.astype(mx.float32)).astype(k_dtype)
+
+        return query, key, value
+
+    @staticmethod
+    def compute_attention(
+        query: mx.array,
+        key: mx.array,
+        value: mx.array,
+        batch_size: int,
+        num_heads: int,
+        head_dim: int,
+        mask: mx.array | None = None,
+    ) -> mx.array:
+        scale = 1 / mx.sqrt(query.shape[-1])
+        hidden_states = scaled_dot_product_attention(query, key, value, scale=scale, mask=mask)
+        hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3))
+        hidden_states = mx.reshape(hidden_states, (batch_size, -1, num_heads * head_dim))
+        return hidden_states
+
+    @staticmethod
+    def convert_key_padding_mask_to_additive_mask(
+        mask: mx.array | None,
+        joint_seq_len: int,
+        txt_seq_len: int,
+    ) -> mx.array | None:
+        if mask is None:
+            return None
+
+        bsz = mask.shape[0]
+        img_seq_len = joint_seq_len - txt_seq_len
+
+        # Create joint mask: [text_mask, image_ones]
+        ones_img = mx.ones((bsz, img_seq_len), dtype=mx.float32)
+        joint_mask = mx.concatenate([mask.astype(mx.float32), ones_img], axis=1)
+
+        # Convert to additive mask for scaled_dot_product_attention
+        additive = (1.0 - joint_mask) * (-1e9)
+        return additive.reshape((additive.shape[0], 1, 1, additive.shape[1]))
+
+    @staticmethod
+    def apply_rope(xq: mx.array, xk: mx.array, freqs_cis: mx.array):
+        xq_ = xq.astype(mx.float32).reshape(*xq.shape[:-1], -1, 1, 2)
+        xk_ = xk.astype(mx.float32).reshape(*xk.shape[:-1], -1, 1, 2)
+        xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+        xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+        return xq_out.reshape(*xq.shape).astype(mx.float32), xk_out.reshape(*xk.shape).astype(mx.float32)
+
+    @staticmethod
+    def apply_rope_bshd(xq: mx.array, xk: mx.array, cos: mx.array, sin: mx.array):
+        xq_f = xq.astype(mx.float32)
+        xk_f = xk.astype(mx.float32)
+        cos_b = cos.reshape(1, cos.shape[0], 1, cos.shape[1])
+        sin_b = sin.reshape(1, sin.shape[0], 1, sin.shape[1])
+
+        def mix(x: mx.array) -> mx.array:
+            x2 = x.reshape(*x.shape[:-1], -1, 2)
+            real = x2[..., 0]
+            imag = x2[..., 1]
+            out0 = real * cos_b + (-imag) * sin_b
+            out1 = imag * cos_b + (real) * sin_b
+            out2 = mx.stack([out0, out1], axis=-1)
+            return out2.reshape(*x.shape)
+
+        return mix(xq_f).astype(mx.float32), mix(xk_f).astype(mx.float32)

src/mflux/models/flux/model/flux_transformer/joint_attention.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from mflux.models.transformer.common.attention_utils import AttentionUtils
+from mflux.models.flux.model.flux_transformer.common.attention_utils import AttentionUtils
 
 
 class JointAttention(nn.Module):

src/mflux/models/flux/model/flux_transformer/joint_transformer_block.py

@@ -1,9 +1,9 @@
 import mlx.core as mx
 from mlx import nn
 
-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
+from mflux.models.flux.model.flux_transformer.ada_layer_norm_zero import AdaLayerNormZero
+from mflux.models.flux.model.flux_transformer.feed_forward import FeedForward
+from mflux.models.flux.model.flux_transformer.joint_attention import JointAttention
 
 
 class JointTransformerBlock(nn.Module):

src/mflux/models/flux/model/flux_transformer/single_block_attention.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 from mlx import nn
 
-from mflux.models.transformer.common.attention_utils import AttentionUtils
+from mflux.models.flux.model.flux_transformer.common.attention_utils import AttentionUtils
 
 
 class SingleBlockAttention(nn.Module):

src/mflux/models/flux/model/flux_transformer/single_transformer_block.py

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

src/mflux/models/flux/model/flux_transformer/time_text_embed.py

@@ -5,9 +5,9 @@ from mlx import nn
 
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
-from mflux.models.transformer.guidance_embedder import GuidanceEmbedder
-from mflux.models.transformer.text_embedder import TextEmbedder
-from mflux.models.transformer.timestep_embedder import TimestepEmbedder
+from mflux.models.flux.model.flux_transformer.guidance_embedder import GuidanceEmbedder
+from mflux.models.flux.model.flux_transformer.text_embedder import TextEmbedder
+from mflux.models.flux.model.flux_transformer.timestep_embedder import TimestepEmbedder
 
 
 class TimeTextEmbed(nn.Module):

src/mflux/models/flux/model/flux_transformer/transformer.py

@@ -5,17 +5,11 @@ from mlx import nn
 
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
-from mflux.models.transformer.ada_layer_norm_continuous 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
+from mflux.models.flux.model.flux_transformer.ada_layer_norm_continuous import AdaLayerNormContinuous
+from mflux.models.flux.model.flux_transformer.embed_nd import EmbedND
+from mflux.models.flux.model.flux_transformer.joint_transformer_block import JointTransformerBlock
+from mflux.models.flux.model.flux_transformer.single_transformer_block import SingleTransformerBlock
+from mflux.models.flux.model.flux_transformer.time_text_embed import TimeTextEmbed
 
 
 class Transformer(nn.Module):
@@ -52,7 +46,7 @@ class Transformer(nn.Module):
         text_embeddings = Transformer.compute_text_embeddings(t, pooled_prompt_embeds, self.time_text_embed, config)
         image_rotary_embeddings = Transformer.compute_rotary_embeddings(prompt_embeds, self.pos_embed, config, kontext_image_ids)  # fmt: off
 
-        # 2. Run the joint transformer blocks
+        # 2. Run the joint flux_transformer blocks
         for idx, block in enumerate(self.transformer_blocks):
             encoder_hidden_states, hidden_states = self._apply_joint_transformer_block(
                 idx=idx,
@@ -67,7 +61,7 @@ class Transformer(nn.Module):
         # 3. Concat the hidden states
         hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
 
-        # 4. Run the single transformer blocks
+        # 4. Run the single flux_transformer blocks
         for idx, block in enumerate(self.single_transformer_blocks):
             hidden_states = self._apply_single_transformer_block(
                 idx=idx,
@@ -95,7 +89,7 @@ class Transformer(nn.Module):
         image_rotary_embeddings: mx.array,
         controlnet_single_block_samples: list[mx.array],
     ) -> mx.array:
-        # 1. Apply single transformer block
+        # 1. Apply single flux_transformer block
         hidden_states = block(
             hidden_states=hidden_states,
             text_embeddings=text_embeddings,
@@ -118,7 +112,7 @@ class Transformer(nn.Module):
         image_rotary_embeddings: mx.array,
         controlnet_block_samples: list[mx.array],
     ) -> mx.array:
-        # 1. Apply joint transformer block
+        # 1. Apply joint flux_transformer block
         encoder_hidden_states, hidden_states = block(
             hidden_states=hidden_states,
             encoder_hidden_states=encoder_hidden_states,

src/mflux/models/flux/model/flux_vae/common/unet_mid_block.py

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

src/mflux/models/flux/model/flux_vae/decoder/decoder.py

@@ -1,13 +1,13 @@
 import mlx.core as mx
 from mlx import nn
 
-from mflux.models.vae.common.unet_mid_block import UnetMidBlock
-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.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
+from mflux.models.flux.model.flux_vae.common.unet_mid_block import UnetMidBlock
+from mflux.models.flux.model.flux_vae.decoder.conv_in import ConvIn
+from mflux.models.flux.model.flux_vae.decoder.conv_norm_out import ConvNormOut
+from mflux.models.flux.model.flux_vae.decoder.conv_out import ConvOut
+from mflux.models.flux.model.flux_vae.decoder.up_block_1_or_2 import UpBlock1Or2
+from mflux.models.flux.model.flux_vae.decoder.up_block_3 import UpBlock3
+from mflux.models.flux.model.flux_vae.decoder.up_block_4 import UpBlock4
 
 
 class Decoder(nn.Module):

src/mflux/models/flux/model/flux_vae/decoder/up_block_1_or_2.py

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