Add after_loop callback to in context generation

src/mflux/community/in_context_lora/flux_in_context_lora.py

@@ -114,7 +114,7 @@ class Flux1InContextLoRA(nn.Module):
                     static_noise=static_noise,
                 )
 
-                # (Optional) Call subscribers at end of loop
+                # (Optional) Call subscribers in-loop
                 Callbacks.in_loop(
                     t=t,
                     seed=seed,
@@ -137,6 +137,14 @@ class Flux1InContextLoRA(nn.Module):
                     time_steps=time_steps,
                 )
 
+        # (Optional) Call subscribers after loop
+        Callbacks.after_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+        )  # fmt: off
+
         # 6. Decode the latent array and return the image
         latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
         decoded = self.vae.decode(latents)

src/mflux/controlnet/flux_controlnet.py

@@ -119,7 +119,7 @@ class Flux1Controlnet(nn.Module):
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
-                # (Optional) Call subscribers at end of loop
+                # (Optional) Call subscribers in-loop
                 Callbacks.in_loop(
                     t=t,
                     seed=seed,
@@ -142,7 +142,7 @@ class Flux1Controlnet(nn.Module):
                     time_steps=time_steps,
                 )
 
-        # (Optional) Call subscribers at end of loop
+        # (Optional) Call subscribers after loop
         Callbacks.after_loop(
             seed=seed,
             prompt=prompt,

src/mflux/flux/flux.py

@@ -99,7 +99,7 @@ class Flux1(nn.Module):
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
-                # (Optional) Call subscribers at end of loop
+                # (Optional) Call subscribers in-loop
                 Callbacks.in_loop(
                     t=t,
                     seed=seed,
@@ -122,7 +122,7 @@ class Flux1(nn.Module):
                     time_steps=time_steps,
                 )
 
-        # (Optional) Call subscribers at end of loop
+        # (Optional) Call subscribers after loop
         Callbacks.after_loop(
             seed=seed,
             prompt=prompt,