Fix imports after empty mflux init file

README.md

@@ -181,7 +181,8 @@ This is useful for integrating MFLUX into shell scripts or dynamically generatin
 Alternatively, you can use MFLUX directly in Python:
 
 ```python
-from mflux import Flux1, Config
+from mflux.flux.flux import Flux1
+from mflux.config.config import Config
 
 # Load the model
 flux = Flux1.from_name(

src/mflux/callbacks/instances/canny_saver.py

@@ -4,9 +4,9 @@ from pathlib import Path
 import mlx.core as mx
 import PIL.Image
 
-from mflux import ImageUtil
 from mflux.callbacks.callback import BeforeLoopCallback
 from mflux.config.runtime_config import RuntimeConfig
+from mflux.post_processing.image_util import ImageUtil
 
 
 class CannyImageSaver(BeforeLoopCallback):

src/mflux/callbacks/instances/depth_saver.py

@@ -4,9 +4,9 @@ from pathlib import Path
 import mlx.core as mx
 import PIL.Image
 
-from mflux import ImageUtil
 from mflux.callbacks.callback import BeforeLoopCallback
 from mflux.config.runtime_config import RuntimeConfig
+from mflux.post_processing.image_util import ImageUtil
 
 
 class DepthImageSaver(BeforeLoopCallback):

src/mflux/concept.py

@@ -1,7 +1,8 @@
-from mflux import Config, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
 from mflux.community.concept_attention.flux_concept import Flux1Concept
-from mflux.error.exceptions import PromptFileReadError
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt

src/mflux/concept_from_image.py

@@ -1,7 +1,8 @@
-from mflux import Config, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
 from mflux.community.concept_attention.flux_concept_from_image import Flux1ConceptFromImage
-from mflux.error.exceptions import PromptFileReadError
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt

src/mflux/controlnet/controlnet_util.py

@@ -21,7 +21,7 @@ class ControlnetUtil:
         controlnet_image_path: StrOrBytesPath,
         is_canny: bool,
     ) -> tuple[mx.array, PIL.Image.Image]:
-        from mflux import ImageUtil
+        from mflux.post_processing.image_util import ImageUtil
 
         control_image = ImageUtil.load_image(controlnet_image_path)
         control_image = ControlnetUtil._scale_image(height=height, width=width, img=control_image)

src/mflux/dreambooth/dataset/dataset.py

@@ -5,11 +5,12 @@ import PIL.Image
 from mlx import nn
 from tqdm import tqdm
 
-from mflux import Flux1, ImageUtil
 from mflux.dreambooth.dataset.batch import Example
 from mflux.dreambooth.dataset.dreambooth_preprocessing import DreamBoothPreProcessing
 from mflux.dreambooth.state.training_spec import ExampleSpec
+from mflux.flux.flux import Flux1
 from mflux.post_processing.array_util import ArrayUtil
+from mflux.post_processing.image_util import ImageUtil
 
 
 class Dataset:

src/mflux/dreambooth/dreambooth.py

@@ -1,12 +1,12 @@
 from mlx import nn
 from tqdm import tqdm
 
-from mflux import Flux1
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.dreambooth.optimization.dreambooth_loss import DreamBoothLoss
 from mflux.dreambooth.state.training_spec import TrainingSpec
 from mflux.dreambooth.state.training_state import TrainingState
 from mflux.dreambooth.statistics.plotter import Plotter
+from mflux.flux.flux import Flux1
 from mflux.weights.weight_handler_lora import WeightHandlerLoRA
 
 

src/mflux/dreambooth/dreambooth_initializer.py

@@ -1,6 +1,7 @@
 import mlx.core.random as random
 
-from mflux import Config, Flux1, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.dreambooth.dataset.dataset import Dataset
 from mflux.dreambooth.dataset.iterator import Iterator
@@ -9,6 +10,7 @@ from mflux.dreambooth.optimization.optimizer import Optimizer
 from mflux.dreambooth.state.training_spec import TrainingSpec
 from mflux.dreambooth.state.training_state import TrainingState
 from mflux.dreambooth.statistics.statistics import Statistics
+from mflux.flux.flux import Flux1
 
 
 class DreamBoothInitializer:

src/mflux/dreambooth/lora_layers/lora_layers.py

@@ -15,7 +15,7 @@ from mflux.utils.version_util import VersionUtil
 from mflux.weights.weight_handler import MetaData, WeightHandler
 
 if TYPE_CHECKING:
-    from mflux import Flux1
+    from mflux.flux.flux import Flux1
 
 
 class LoRALayers:

src/mflux/dreambooth/optimization/dreambooth_loss.py

@@ -2,10 +2,11 @@ import random
 
 import mlx.core as mx
 
-from mflux import Config, Flux1
+from mflux.config.config import Config
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.dreambooth.dataset.batch import Batch
 from mflux.dreambooth.dataset.dataset import Example
+from mflux.flux.flux import Flux1
 from mflux.latent_creator.latent_creator import LatentCreator
 
 

src/mflux/flux/flux_initializer.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+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

src/mflux/flux_tools/redux/redux_util.py

@@ -2,9 +2,9 @@ from pathlib import Path
 
 import mlx.core as mx
 
-from mflux import ImageUtil
 from mflux.models.redux_encoder.redux_encoder import ReduxEncoder
 from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer
+from mflux.post_processing.image_util import ImageUtil
 
 
 class ReduxUtil:

src/mflux/flux_tools/redux/weight_handler_redux.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from mflux.weights.download import snapshot_download
 from mflux.weights.weight_handler import MetaData, WeightHandler
 

src/mflux/generate.py

@@ -1,6 +1,8 @@
-from mflux import Config, Flux1, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.error.exceptions import PromptFileReadError
+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.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt

src/mflux/generate_controlnet.py

@@ -1,6 +1,8 @@
-from mflux import Config, Flux1Controlnet, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.error.exceptions import PromptFileReadError
+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.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt

src/mflux/generate_depth.py

@@ -1,6 +1,6 @@
-from mflux import Config, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.error.exceptions import PromptFileReadError
+from mflux.config.config import Config
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.flux_tools.depth.flux_depth import Flux1Depth
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser

src/mflux/generate_fill.py

@@ -1,6 +1,6 @@
-from mflux import Config, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.error.exceptions import PromptFileReadError
+from mflux.config.config import Config
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.flux_tools.fill.flux_fill import Flux1Fill
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser

src/mflux/generate_in_context_catvton.py

@@ -1,9 +1,10 @@
 from pathlib import Path
 
-from mflux import Config, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
 from mflux.community.in_context.flux_in_context_fill import Flux1InContextFill
-from mflux.error.exceptions import PromptFileReadError
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt

src/mflux/generate_in_context_dev.py

@@ -1,11 +1,11 @@
 from pathlib import Path
 
-from mflux import Config, StopImageGenerationException
 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
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt

src/mflux/generate_in_context_edit.py

@@ -2,11 +2,12 @@ from pathlib import Path
 
 from PIL import Image
 
-from mflux import Config, ModelConfig, StopImageGenerationException
 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.error.exceptions import PromptFileReadError
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt

src/mflux/generate_kontext.py

@@ -1,8 +1,8 @@
 from pathlib import Path
 
-from mflux import Config, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.error.exceptions import PromptFileReadError
+from mflux.config.config import Config
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.kontext.flux_kontext import Flux1Kontext
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser

src/mflux/generate_redux.py

@@ -1,8 +1,9 @@
 from pathlib import Path
 
-from mflux import Config, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.error.exceptions import PromptFileReadError
+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.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser

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

@@ -2,7 +2,7 @@ import mlx.core as mx
 from mlx import nn
 from mlx.core.fast import scaled_dot_product_attention
 
-from mflux import Config
+from mflux.config.config import Config
 
 
 class CLIPSdpaAttention(nn.Module):

src/mflux/post_processing/generated_image.py

@@ -86,7 +86,7 @@ class GeneratedImage:
         export_json_metadata: bool = False,
         overwrite: bool = False,
     ) -> None:
-        from mflux import ImageUtil
+        from mflux.post_processing.image_util import ImageUtil
 
         ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata, overwrite)
 
@@ -112,7 +112,7 @@ class GeneratedImage:
         overwrite: bool = False,
     ) -> None:
         if self.concept_heatmap:
-            from mflux import ImageUtil
+            from mflux.post_processing.image_util import ImageUtil
 
             ImageUtil.save_image(
                 image=self.concept_heatmap.image,

src/mflux/save.py

@@ -1,4 +1,5 @@
-from mflux import Flux1, ModelConfig
+from mflux.config.model_config import ModelConfig
+from mflux.flux.flux import Flux1
 from mflux.ui.cli.parsers import CommandLineParser
 
 

src/mflux/upscale.py

@@ -2,9 +2,11 @@ import sys
 
 import PIL.Image
 
-from mflux import Config, Flux1Controlnet, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_manager import CallbackManager
-from mflux.error.exceptions import PromptFileReadError
+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.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_prompt

tests/dreambooth/test_train_and_load_weights.py

@@ -3,10 +3,12 @@ import shutil
 
 import numpy as np
 
-from mflux import Config, Flux1, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 from mflux.dreambooth.dreambooth import DreamBooth
 from mflux.dreambooth.dreambooth_initializer import DreamBoothInitializer
 from mflux.dreambooth.state.zip_util import ZipUtil
+from mflux.flux.flux import Flux1
 
 CHECKPOINT = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint.zip"
 OUTPUT_DIR = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint"

tests/image_generation/helpers/image_generation_concept_test_helper.py

@@ -4,9 +4,10 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
-from mflux import Config, ModelConfig
 from mflux.community.concept_attention.flux_concept import Flux1Concept
 from mflux.community.concept_attention.flux_concept_from_image import Flux1ConceptFromImage
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 
 
 class ImageGenerationConceptTestHelper:

tests/image_generation/helpers/image_generation_controlnet_test_helper.py

@@ -3,7 +3,9 @@ import os
 import numpy as np
 from PIL import Image
 
-from mflux import Config, Flux1Controlnet, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.controlnet.flux_controlnet import Flux1Controlnet
 from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
 
 

tests/image_generation/helpers/image_generation_depth_test_helper.py

@@ -4,7 +4,8 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
-from mflux import Config, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 from mflux.flux_tools.depth.flux_depth import Flux1Depth
 
 

tests/image_generation/helpers/image_generation_fill_test_helper.py

@@ -3,7 +3,8 @@ import os
 import numpy as np
 from PIL import Image
 
-from mflux import Config, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 from mflux.flux_tools.fill.flux_fill import Flux1Fill
 from mflux.ui import defaults as ui_defaults
 from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper

tests/image_generation/helpers/image_generation_ic_edit_test_helper.py

@@ -4,9 +4,10 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
-from mflux import Config, ModelConfig
 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
 
 
 class ImageGeneratorICEditTestHelper:

tests/image_generation/helpers/image_generation_in_context_test_helper.py

@@ -4,9 +4,10 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
-from mflux import Config, ModelConfig
 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
 
 
 class ImageGeneratorInContextTestHelper:

tests/image_generation/helpers/image_generation_kontext_test_helper.py

@@ -4,7 +4,8 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
-from mflux import Config, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 from mflux.kontext.flux_kontext import Flux1Kontext
 
 

tests/image_generation/helpers/image_generation_redux_test_helper.py

@@ -4,7 +4,8 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
-from mflux import Config, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 from mflux.flux_tools.redux.flux_redux import Flux1Redux
 
 

tests/image_generation/helpers/image_generation_test_helper.py

@@ -4,7 +4,9 @@ from pathlib import Path
 import numpy as np
 from PIL import Image
 
-from mflux import Config, Flux1, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.flux.flux import Flux1
 
 
 class ImageGeneratorTestHelper:

tests/image_generation/test_generate_concept.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from tests.image_generation.helpers.image_generation_concept_test_helper import ImageGenerationConceptTestHelper
 
 

tests/image_generation/test_generate_image.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
 
 

tests/image_generation/test_generate_image_controlnet.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from tests.image_generation.helpers.image_generation_controlnet_test_helper import ImageGeneratorControlnetTestHelper
 
 

tests/image_generation/test_generate_image_depth.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from tests.image_generation.helpers.image_generation_depth_test_helper import ImageGeneratorDepthTestHelper
 
 

tests/image_generation/test_generate_image_fill.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from tests.image_generation.helpers.image_generation_fill_test_helper import ImageGeneratorFillTestHelper
 
 

tests/image_generation/test_generate_image_in_context.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from tests.image_generation.helpers.image_generation_ic_edit_test_helper import ImageGeneratorICEditTestHelper
 from tests.image_generation.helpers.image_generation_in_context_test_helper import ImageGeneratorInContextTestHelper
 

tests/image_generation/test_generate_image_kontext.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from tests.image_generation.helpers.image_generation_kontext_test_helper import ImageGeneratorKontextTestHelper
 
 

tests/image_generation/test_generate_image_redux.py

@@ -1,4 +1,4 @@
-from mflux import ModelConfig
+from mflux.config.model_config import ModelConfig
 from tests.image_generation.helpers.image_generation_redux_test_helper import ImageGeneratorReduxTestHelper
 
 

tests/model_saving/test_model_saving.py

@@ -4,7 +4,9 @@ from pathlib import Path
 
 import numpy as np
 
-from mflux import Config, Flux1, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.flux.flux import Flux1
 from mflux.utils.version_util import VersionUtil
 from mflux.weights.weight_handler import WeightHandler
 

tests/model_saving/test_model_saving_lora.py

@@ -4,7 +4,9 @@ from pathlib import Path
 
 import numpy as np
 
-from mflux import Config, Flux1, ModelConfig
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.flux.flux import Flux1
 
 PATH = "tests/4bit/"