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

Merge pull request #126 from filipstrand/img2img-special-case

Browse Source 
Img2img special case fix
This commit is contained in:
Filip Strand 2025-02-14 10:41:02 +01:00 committed by GitHub
commit dca6c0acea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 92 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/mflux/callbacks/callback.py
View File

@ -12,6 +12,8 @@ class BeforeLoopCallback(Protocol):
self,
seed: int,
prompt: str,
latents: mx.array,
config: RuntimeConfig,
canny_image: PIL.Image.Image | None = None,
) -> None: # fmt: off
...
@ -20,9 +22,9 @@ class BeforeLoopCallback(Protocol):
class InLoopCallback(Protocol):
def call_in_loop(
self,
t: int,
seed: int,
prompt: str,
step: int,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm
@ -33,9 +35,9 @@ class InLoopCallback(Protocol):
class InterruptCallback(Protocol):
def call_interrupt(
self,
t: int,
seed: int,
prompt: str,
step: int,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm

12
src/mflux/callbacks/callbacks.py
View File

@ -11,29 +11,33 @@ class Callbacks:
def before_loop(
seed: int,
prompt: str,
latents: mx.array,
config: RuntimeConfig,
canny_image: PIL.Image.Image | None = None,
): # fmt: off
for subscriber in CallbackRegistry.before_loop_callbacks():
subscriber.call_before_loop(
seed=seed,
prompt=prompt,
latents=latents,
config=config,
canny_image=canny_image,
)
@staticmethod
def in_loop(
t: int,
seed: int,
prompt: str,
step: int,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm
): # fmt: off
for subscriber in CallbackRegistry.in_loop_callbacks():
subscriber.call_in_loop(
t=t,
seed=seed,
prompt=prompt,
step=step,
latents=latents,
config=config,
time_steps=time_steps,
@ -41,18 +45,18 @@ class Callbacks:
@staticmethod
def interruption(
t: int,
seed: int,
prompt: str,
step: int,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm
): # fmt: off
for subscriber in CallbackRegistry.interrupt_callbacks():
subscriber.call_interrupt(
t=t,
seed=seed,
prompt=prompt,
step=step,
latents=latents,
config=config,
time_steps=time_steps,

6
src/mflux/callbacks/instances/canny_saver.py
View File

@ -1,10 +1,12 @@
import os
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
class CannyImageSaver(BeforeLoopCallback):
@ -15,8 +17,10 @@ class CannyImageSaver(BeforeLoopCallback):
self,
seed: int,
prompt: str,
latents: mx.array,
config: RuntimeConfig,
canny_image: PIL.Image.Image | None = None,
) -> None: # fmt: off
) -> None:
base, ext = os.path.splitext(self.path)
ImageUtil.save_image(
image=canny_image,

74
src/mflux/callbacks/instances/stepwise_handler.py
View File

@ -1,16 +1,17 @@
from pathlib import Path
import mlx.core as mx
import PIL.Image
import tqdm
from mflux import StopImageGenerationException
from mflux.callbacks.callback import InLoopCallback, InterruptCallback
from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback
from mflux.config.runtime_config import RuntimeConfig
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.image_util import ImageUtil
class StepwiseHandler(InLoopCallback, InterruptCallback):
class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback):
def __init__(
self,
flux,
@ -23,17 +24,65 @@ class StepwiseHandler(InLoopCallback, InterruptCallback):
if self.output_dir:
self.output_dir.mkdir(parents=True, exist_ok=True)
def call_in_loop(
def call_before_loop(
self,
seed: int,
prompt: str,
latents: mx.array,
config: RuntimeConfig,
canny_image: PIL.Image.Image | None = None,
) -> None: # fmt: off
self._save_image(
step=config.init_time_step,
seed=seed,
prompt=prompt,
latents=latents,
config=config,
time_steps=None,
)
def call_in_loop(
self,
t: int,
seed: int,
prompt: str,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm
) -> None: # fmt: off
self._save_image(
step=t + 1,
seed=seed,
prompt=prompt,
latents=latents,
config=config,
time_steps=time_steps,
)
def call_interrupt(
self,
t: int,
seed: int,
prompt: str,
latents: mx.array,
config: RuntimeConfig,
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,
step: int,
seed: int,
prompt: str,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm
) -> None: # fmt: off
unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
stepwise_decoded = self.flux.vae.decode(unpack_latents)
generation_time = time_steps.format_dict["elapsed"] if time_steps is not None else 0
stepwise_img = ImageUtil.to_image(
decoded_latents=stepwise_decoded,
config=config,
@ -42,28 +91,15 @@ class StepwiseHandler(InLoopCallback, InterruptCallback):
quantization=self.flux.bits,
lora_paths=self.flux.lora_paths,
lora_scales=self.flux.lora_scales,
generation_time=time_steps.format_dict["elapsed"],
generation_time=generation_time,
)
self.step_wise_images.append(stepwise_img)
stepwise_img.save(
path=self.output_dir / f"seed_{seed}_step{step}of{len(time_steps)}.png",
path=self.output_dir / f"seed_{seed}_step{step}of{config.num_inference_steps}.png",
export_json_metadata=False,
)
self.step_wise_images.append(stepwise_img)
self._save_composite(seed=seed)
def call_interrupt(
self,
seed: int,
prompt: str,
step: int,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm
) -> None: # fmt: off
self._save_composite(seed=seed)
raise StopImageGenerationException(f"Stopping image generation at step {step + 1}/{len(time_steps)}")
def _save_composite(self, seed: int) -> None:
if self.step_wise_images:
composite_img = ImageUtil.to_composite_image(self.step_wise_images)

12
src/mflux/config/runtime_config.py
View File

@ -53,14 +53,16 @@ class RuntimeConfig:
@property
def init_time_step(self) -> int:
if self.config.init_image_path is None:
# For text-to-image, always begin at time step 0.
is_txt2img = self.config.init_image_path is None or self.config.init_image_strength == 0.0
if is_txt2img:
return 0
else:
# For image-to-image: higher strength means we skip more steps.
# 1. Clamp strength to [0, 1]
strength = max(0.0, min(1.0, self.config.init_image_strength))
t = max(1, int(self.num_inference_steps * strength))
return t
# 2. Return start time in [1, floor(num_steps * strength)]
return max(1, int(self.num_inference_steps * strength))
@property
def controlnet_strength(self) -> float | None:

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

@ -87,10 +87,12 @@ class Flux1Controlnet(nn.Module):
Callbacks.before_loop(
seed=seed,
prompt=prompt,
latents=latents,
config=config,
canny_image=canny_image
) # fmt: off
for gen_step, t in enumerate(time_steps, 1):
for t in time_steps:
try:
# 4.t Compute controlnet samples
controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet(
@ -119,9 +121,9 @@ class Flux1Controlnet(nn.Module):
# (Optional) Call subscribes at end of loop
Callbacks.in_loop(
t=t,
seed=seed,
prompt=prompt,
step=gen_step,
latents=latents,
config=config,
time_steps=time_steps,
@ -132,9 +134,9 @@ class Flux1Controlnet(nn.Module):
except KeyboardInterrupt: # noqa: PERF203
Callbacks.interruption(
t=t,
seed=seed,
prompt=prompt,
step=gen_step,
latents=latents,
config=config,
time_steps=time_steps,

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

@ -79,10 +79,12 @@ class Flux1(nn.Module):
# (Optional) Call subscribers for beginning of loop
Callbacks.before_loop(
seed=seed,
prompt=prompt
prompt=prompt,
latents=latents,
config=config,
) # fmt: off
for gen_step, t in enumerate(time_steps, 1):
for t in time_steps:
try:
# 3.t Predict the noise
noise = self.transformer(
@ -99,9 +101,9 @@ class Flux1(nn.Module):
# (Optional) Call subscribes at end of loop
Callbacks.in_loop(
t=t,
seed=seed,
prompt=prompt,
step=gen_step,
latents=latents,
config=config,
time_steps=time_steps,
@ -112,9 +114,9 @@ class Flux1(nn.Module):
except KeyboardInterrupt: # noqa: PERF203
Callbacks.interruption(
t=t,
seed=seed,
prompt=prompt,
step=gen_step,
latents=latents,
config=config,
time_steps=time_steps,

1
src/mflux/generate.py
View File

@ -26,6 +26,7 @@ def main():
# 2. Register the optional callbacks
if args.stepwise_image_output_dir:
handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
CallbackRegistry.register_before_loop(handler)
CallbackRegistry.register_in_loop(handler)
CallbackRegistry.register_interrupt(handler)

1
src/mflux/generate_controlnet.py
View File

@ -29,6 +29,7 @@ def main():
CallbackRegistry.register_before_loop(CannyImageSaver(path=args.output))
if args.stepwise_image_output_dir:
handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
CallbackRegistry.register_before_loop(handler)
CallbackRegistry.register_in_loop(handler)
CallbackRegistry.register_interrupt(handler)