Merge branch 'main' into prepare-outpaint-scaffolding

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mflux"
-version = "0.6.0"
+version = "0.6.2"
 description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
 readme = "README.md"
 keywords = ["diffusers", "flux", "mlx"]

src/mflux/callbacks/instances/stepwise_handler.py

@@ -4,7 +4,6 @@ import mlx.core as mx
 import PIL.Image
 import tqdm
 
-from mflux import StopImageGenerationException
 from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.post_processing.array_util import ArrayUtil
@@ -69,7 +68,6 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback):
         time_steps: tqdm
     ) -> None:  # fmt: off
         self._save_composite(seed=seed)
-        raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 
     def _save_image(
         self,

src/mflux/community/in_context_lora/flux_in_context_lora.py

@@ -6,6 +6,7 @@ from mflux.callbacks.callbacks import Callbacks
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
+from mflux.error.exceptions import StopImageGenerationException
 from mflux.flux.flux_initializer import FluxInitializer
 from mflux.latent_creator.latent_creator import Img2Img, LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
@@ -136,6 +137,7 @@ class Flux1InContextLoRA(nn.Module):
                     config=config,
                     time_steps=time_steps,
                 )
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 
         # (Optional) Call subscribers after loop
         Callbacks.after_loop(

src/mflux/controlnet/flux_controlnet.py

@@ -8,6 +8,7 @@ from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.controlnet.controlnet_util import ControlnetUtil
 from mflux.controlnet.transformer_controlnet import TransformerControlnet
+from mflux.error.exceptions import StopImageGenerationException
 from mflux.flux.flux_initializer import FluxInitializer
 from mflux.latent_creator.latent_creator import LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
@@ -141,6 +142,7 @@ class Flux1Controlnet(nn.Module):
                     config=config,
                     time_steps=time_steps,
                 )
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 
         # (Optional) Call subscribers after loop
         Callbacks.after_loop(

src/mflux/flux/flux.py

@@ -6,6 +6,7 @@ from mflux.callbacks.callbacks import Callbacks
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
+from mflux.error.exceptions import StopImageGenerationException
 from mflux.flux.flux_initializer import FluxInitializer
 from mflux.latent_creator.latent_creator import Img2Img, LatentCreator
 from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
@@ -121,6 +122,7 @@ class Flux1(nn.Module):
                     config=config,
                     time_steps=time_steps,
                 )
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
 
         # (Optional) Call subscribers after loop
         Callbacks.after_loop(

src/mflux/ui/cli/parsers.py

@@ -214,7 +214,7 @@ class CommandLineParser(argparse.ArgumentParser):
             namespace.image_outpaint_padding = box_values.parse_box_value(namespace.image_outpaint_padding)
             print(f"{namespace.image_outpaint_padding=}")
 
-        if hasattr(namespace, "low_ram") and namespace.low_ram is not None and len(namespace.seed) > 1:
+        if getattr(namespace, 'low_ram', False) and len(namespace.seed) > 1:
             self.error("--low-ram cannot be used with multiple seeds")
 
         return namespace

tests/arg_parser/test_cli_argparser.py

@@ -232,7 +232,7 @@ def test_auto_seeds_arg(mflux_generate_parser, mflux_generate_minimal_model_argv
         assert "_seed_{seed}" in args.output
 
     for _ in range(0, 10):
-        random_auto_seed_count = random.randint(0, 100)
+        random_auto_seed_count = random.randint(2, 100)
         with patch("sys.argv", mflux_generate_minimal_model_argv + ["--auto-seeds", str(random_auto_seed_count)]):
             args = mflux_generate_parser.parse_args()
             assert len(set(args.seed)) == random_auto_seed_count