diff --git a/README.md b/README.md index bd61ca5..3590a2b 100644 --- a/README.md +++ b/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( diff --git a/src/mflux/callbacks/instances/canny_saver.py b/src/mflux/callbacks/instances/canny_saver.py index 8cb3850..5b9e364 100644 --- a/src/mflux/callbacks/instances/canny_saver.py +++ b/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): diff --git a/src/mflux/callbacks/instances/depth_saver.py b/src/mflux/callbacks/instances/depth_saver.py index 0d5a561..e07bd9d 100644 --- a/src/mflux/callbacks/instances/depth_saver.py +++ b/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): diff --git a/src/mflux/concept.py b/src/mflux/concept.py index c3bfa16..a62223d 100644 --- a/src/mflux/concept.py +++ b/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 diff --git a/src/mflux/concept_from_image.py b/src/mflux/concept_from_image.py index f5841be..f7eef50 100644 --- a/src/mflux/concept_from_image.py +++ b/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 diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py index 5bd3330..da38d70 100644 --- a/src/mflux/controlnet/controlnet_util.py +++ b/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) diff --git a/src/mflux/dreambooth/dataset/dataset.py b/src/mflux/dreambooth/dataset/dataset.py index 0c386a1..82d3bf4 100644 --- a/src/mflux/dreambooth/dataset/dataset.py +++ b/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: diff --git a/src/mflux/dreambooth/dreambooth.py b/src/mflux/dreambooth/dreambooth.py index b6ad86e..16bf853 100644 --- a/src/mflux/dreambooth/dreambooth.py +++ b/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 diff --git a/src/mflux/dreambooth/dreambooth_initializer.py b/src/mflux/dreambooth/dreambooth_initializer.py index b560da6..629a969 100644 --- a/src/mflux/dreambooth/dreambooth_initializer.py +++ b/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: diff --git a/src/mflux/dreambooth/lora_layers/lora_layers.py b/src/mflux/dreambooth/lora_layers/lora_layers.py index 16039a6..f794d4a 100644 --- a/src/mflux/dreambooth/lora_layers/lora_layers.py +++ b/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: diff --git a/src/mflux/dreambooth/optimization/dreambooth_loss.py b/src/mflux/dreambooth/optimization/dreambooth_loss.py index 1740cfa..6014ae3 100644 --- a/src/mflux/dreambooth/optimization/dreambooth_loss.py +++ b/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 diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index 1ae87f5..7746b34 100644 --- a/src/mflux/flux/flux_initializer.py +++ b/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 diff --git a/src/mflux/flux_tools/redux/redux_util.py b/src/mflux/flux_tools/redux/redux_util.py index 01cb57f..0d3e27b 100644 --- a/src/mflux/flux_tools/redux/redux_util.py +++ b/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: diff --git a/src/mflux/flux_tools/redux/weight_handler_redux.py b/src/mflux/flux_tools/redux/weight_handler_redux.py index 6140c0f..05e3ca5 100644 --- a/src/mflux/flux_tools/redux/weight_handler_redux.py +++ b/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 diff --git a/src/mflux/generate.py b/src/mflux/generate.py index ac75322..f2df4f7 100644 --- a/src/mflux/generate.py +++ b/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 diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 7c843d1..d03f0a1 100644 --- a/src/mflux/generate_controlnet.py +++ b/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 diff --git a/src/mflux/generate_depth.py b/src/mflux/generate_depth.py index 1a5d30a..203ddcf 100644 --- a/src/mflux/generate_depth.py +++ b/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 diff --git a/src/mflux/generate_fill.py b/src/mflux/generate_fill.py index 7e34ae9..ba61d6b 100644 --- a/src/mflux/generate_fill.py +++ b/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 diff --git a/src/mflux/generate_in_context_catvton.py b/src/mflux/generate_in_context_catvton.py index e7d5d9f..cf38c23 100644 --- a/src/mflux/generate_in_context_catvton.py +++ b/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 diff --git a/src/mflux/generate_in_context_dev.py b/src/mflux/generate_in_context_dev.py index 9e1c778..40dfab7 100644 --- a/src/mflux/generate_in_context_dev.py +++ b/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 diff --git a/src/mflux/generate_in_context_edit.py b/src/mflux/generate_in_context_edit.py index 062c816..7ce958e 100644 --- a/src/mflux/generate_in_context_edit.py +++ b/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 diff --git a/src/mflux/generate_kontext.py b/src/mflux/generate_kontext.py index 6b43eec..ca27ce4 100644 --- a/src/mflux/generate_kontext.py +++ b/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 diff --git a/src/mflux/generate_redux.py b/src/mflux/generate_redux.py index 86ce345..09c5512 100644 --- a/src/mflux/generate_redux.py +++ b/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 diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py b/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py index e067d1d..cc560ea 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py +++ b/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): diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index 225c038..91fb352 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/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, diff --git a/src/mflux/save.py b/src/mflux/save.py index cd96b98..6fd0e81 100644 --- a/src/mflux/save.py +++ b/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 diff --git a/src/mflux/upscale.py b/src/mflux/upscale.py index 226177c..a28fa81 100644 --- a/src/mflux/upscale.py +++ b/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 diff --git a/tests/dreambooth/test_train_and_load_weights.py b/tests/dreambooth/test_train_and_load_weights.py index 0d0697e..1b4b30c 100644 --- a/tests/dreambooth/test_train_and_load_weights.py +++ b/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" diff --git a/tests/image_generation/helpers/image_generation_concept_test_helper.py b/tests/image_generation/helpers/image_generation_concept_test_helper.py index 64906d9..48d0102 100644 --- a/tests/image_generation/helpers/image_generation_concept_test_helper.py +++ b/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: diff --git a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py index 6a8f0b0..fbf6009 100644 --- a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py +++ b/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 diff --git a/tests/image_generation/helpers/image_generation_depth_test_helper.py b/tests/image_generation/helpers/image_generation_depth_test_helper.py index af275d4..9487e1e 100644 --- a/tests/image_generation/helpers/image_generation_depth_test_helper.py +++ b/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 diff --git a/tests/image_generation/helpers/image_generation_fill_test_helper.py b/tests/image_generation/helpers/image_generation_fill_test_helper.py index 09d6ea3..5a6a914 100644 --- a/tests/image_generation/helpers/image_generation_fill_test_helper.py +++ b/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 diff --git a/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py b/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py index b0f0ba1..62a4afe 100644 --- a/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py +++ b/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: diff --git a/tests/image_generation/helpers/image_generation_in_context_test_helper.py b/tests/image_generation/helpers/image_generation_in_context_test_helper.py index 7c99c78..f912564 100644 --- a/tests/image_generation/helpers/image_generation_in_context_test_helper.py +++ b/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: diff --git a/tests/image_generation/helpers/image_generation_kontext_test_helper.py b/tests/image_generation/helpers/image_generation_kontext_test_helper.py index 489efc3..acd4e60 100644 --- a/tests/image_generation/helpers/image_generation_kontext_test_helper.py +++ b/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 diff --git a/tests/image_generation/helpers/image_generation_redux_test_helper.py b/tests/image_generation/helpers/image_generation_redux_test_helper.py index 3da822d..5dc68b4 100644 --- a/tests/image_generation/helpers/image_generation_redux_test_helper.py +++ b/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 diff --git a/tests/image_generation/helpers/image_generation_test_helper.py b/tests/image_generation/helpers/image_generation_test_helper.py index 414be74..ac63ba6 100644 --- a/tests/image_generation/helpers/image_generation_test_helper.py +++ b/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: diff --git a/tests/image_generation/test_generate_concept.py b/tests/image_generation/test_generate_concept.py index 4d186e0..e7d668e 100644 --- a/tests/image_generation/test_generate_concept.py +++ b/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 diff --git a/tests/image_generation/test_generate_image.py b/tests/image_generation/test_generate_image.py index 4662192..89b8eab 100644 --- a/tests/image_generation/test_generate_image.py +++ b/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 diff --git a/tests/image_generation/test_generate_image_controlnet.py b/tests/image_generation/test_generate_image_controlnet.py index 0eb7d50..311e168 100644 --- a/tests/image_generation/test_generate_image_controlnet.py +++ b/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 diff --git a/tests/image_generation/test_generate_image_depth.py b/tests/image_generation/test_generate_image_depth.py index 6e9a612..a4c0842 100644 --- a/tests/image_generation/test_generate_image_depth.py +++ b/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 diff --git a/tests/image_generation/test_generate_image_fill.py b/tests/image_generation/test_generate_image_fill.py index cba1984..9b849c8 100644 --- a/tests/image_generation/test_generate_image_fill.py +++ b/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 diff --git a/tests/image_generation/test_generate_image_in_context.py b/tests/image_generation/test_generate_image_in_context.py index 0bccd81..0c0684f 100644 --- a/tests/image_generation/test_generate_image_in_context.py +++ b/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 diff --git a/tests/image_generation/test_generate_image_kontext.py b/tests/image_generation/test_generate_image_kontext.py index 46ecf63..71556a1 100644 --- a/tests/image_generation/test_generate_image_kontext.py +++ b/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 diff --git a/tests/image_generation/test_generate_image_redux.py b/tests/image_generation/test_generate_image_redux.py index fde4c7f..f1b5fa6 100644 --- a/tests/image_generation/test_generate_image_redux.py +++ b/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 diff --git a/tests/model_saving/test_model_saving.py b/tests/model_saving/test_model_saving.py index 842b6bb..6f7f7e0 100644 --- a/tests/model_saving/test_model_saving.py +++ b/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 diff --git a/tests/model_saving/test_model_saving_lora.py b/tests/model_saving/test_model_saving_lora.py index 88bf8ab..3851f81 100644 --- a/tests/model_saving/test_model_saving_lora.py +++ b/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/"