monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add after_loop callback to in context generation

Browse Source
This commit is contained in:
filipstrand 2025-03-05 08:29:49 +01:00
parent 368fa0faf7
commit e13181e176
3 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
src/mflux/community/in_context_lora/flux_in_context_lora.py
View File

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

4
src/mflux/controlnet/flux_controlnet.py
View File

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

4
src/mflux/flux/flux.py
View File

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