From 04d24234f104ace5460457c4b0143ef8d564f5c3 Mon Sep 17 00:00:00 2001
From: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>
Date: Sun, 9 Mar 2025 07:42:12 -0700
Subject: [PATCH] Fix: image gen should stop even if no callbacks are
 registered

---
 pyproject.toml                                              | 2 +-
 src/mflux/callbacks/instances/stepwise_handler.py           | 2 --
 src/mflux/community/in_context_lora/flux_in_context_lora.py | 2 ++
 src/mflux/controlnet/flux_controlnet.py                     | 2 ++
 src/mflux/flux/flux.py                                      | 2 ++
 5 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index d79f3dd..7e7baa6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mflux"
-version = "0.6.0"
+version = "0.6.1"
 description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
 readme = "README.md"
 keywords = ["diffusers", "flux", "mlx"]
diff --git a/src/mflux/callbacks/instances/stepwise_handler.py b/src/mflux/callbacks/instances/stepwise_handler.py
index ea74c60..8e0f0ba 100644
--- a/src/mflux/callbacks/instances/stepwise_handler.py
+++ b/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,
diff --git a/src/mflux/community/in_context_lora/flux_in_context_lora.py b/src/mflux/community/in_context_lora/flux_in_context_lora.py
index 9c27fc2..39ca300 100644
--- a/src/mflux/community/in_context_lora/flux_in_context_lora.py
+++ b/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(
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 5c9d878..03b7401 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/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(
diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py
index bf7f756..a0fa92f 100644
--- a/src/mflux/flux/flux.py
+++ b/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(