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

Implement general callback mechanism

Browse Source
This commit is contained in:
filipstrand 2025-02-09 11:29:50 +01:00
parent e4af40e4ee
commit 3601baf75e
15 changed files with 385 additions and 186 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

0
src/mflux/callbacks/__init__.py Normal file
View File

43
src/mflux/callbacks/callback.py Normal file
View File

@ -0,0 +1,43 @@
from typing import Protocol
import mlx.core as mx
import PIL.Image
import tqdm
from mflux.config.runtime_config import RuntimeConfig
class BeforeLoopCallback(Protocol):
def call_before_loop(
self,
seed: int,
prompt: str,
canny_image: PIL.Image.Image | None = None,
) -> None: # fmt: off
...
class InLoopCallback(Protocol):
def call_in_loop(
self,
seed: int,
prompt: str,
step: int,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm
) -> None: # fmt: off
...
class InterruptCallback(Protocol):
def call_interrupt(
self,
seed: int,
prompt: str,
step: int,
latents: mx.array,
config: RuntimeConfig,
time_steps: tqdm
) -> None: # fmt: off
...

31
src/mflux/callbacks/callback_registry.py Normal file
View File

@ -0,0 +1,31 @@
from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback
class CallbackRegistry:
in_loop = []
before_loop = []
interrupt = []
@staticmethod
def register_in_loop(callback: InLoopCallback) -> None:
CallbackRegistry.in_loop.append(callback)
@staticmethod
def register_before_loop(callback: BeforeLoopCallback) -> None:
CallbackRegistry.before_loop.append(callback)
@staticmethod
def register_interrupt(callback: InterruptCallback) -> None:
CallbackRegistry.interrupt.append(callback)
@staticmethod
def in_loop_callbacks() -> list[InLoopCallback]:
return CallbackRegistry.in_loop
@staticmethod
def before_loop_callbacks() -> list[BeforeLoopCallback]:
return CallbackRegistry.before_loop
@staticmethod
def interrupt_callbacks() -> list[InterruptCallback]:
return CallbackRegistry.interrupt

59
src/mflux/callbacks/callbacks.py Normal file
View File

@ -0,0 +1,59 @@
import mlx.core as mx
import PIL.Image
import tqdm
from mflux.callbacks.callback_registry import CallbackRegistry
from mflux.config.runtime_config import RuntimeConfig
class Callbacks:
@staticmethod
def before_loop(
seed: int,
prompt: str,
canny_image: PIL.Image.Image | None = None,
): # fmt: off
for subscriber in CallbackRegistry.before_loop_callbacks():
subscriber.call_before_loop(
seed=seed,
prompt=prompt,
canny_image=canny_image,
)
@staticmethod
def in_loop(
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(
seed=seed,
prompt=prompt,
step=step,
latents=latents,
config=config,
time_steps=time_steps,
)
@staticmethod
def interruption(
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(
seed=seed,
prompt=prompt,
step=step,
latents=latents,
config=config,
time_steps=time_steps,
)

0
src/mflux/callbacks/instances/__init__.py Normal file
View File

24
src/mflux/callbacks/instances/canny_saver.py Normal file
View File

@ -0,0 +1,24 @@
import os
from pathlib import Path
import PIL.Image
from mflux import ImageUtil
from mflux.callbacks.callback import BeforeLoopCallback
class CannyImageSaver(BeforeLoopCallback):
def __init__(self, path: str):
self.path = Path(path)
def call_before_loop(
self,
seed: int,
prompt: str,
canny_image: PIL.Image.Image | None = None,
) -> None: # fmt: off
base, ext = os.path.splitext(self.path)
ImageUtil.save_image(
image=canny_image,
path=f"{base}_controlnet_canny{ext}"
) # fmt: off

70
src/mflux/callbacks/instances/stepwise_handler.py Normal file
View File

@ -0,0 +1,70 @@
from pathlib import Path
import mlx.core as mx
import tqdm
from mflux import StopImageGenerationException
from mflux.callbacks.callback import 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):
def __init__(
self,
flux,
output_dir: str,
):
self.flux = flux
self.output_dir = Path(output_dir)
self.step_wise_images = []
if self.output_dir:
self.output_dir.mkdir(parents=True, exist_ok=True)
def call_in_loop(
self,
seed: int,
prompt: str,
step: int,
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)
stepwise_img = ImageUtil.to_image(
decoded_latents=stepwise_decoded,
config=config,
seed=seed,
prompt=prompt,
quantization=self.flux.bits,
lora_paths=self.flux.lora_paths,
lora_scales=self.flux.lora_scales,
generation_time=time_steps.format_dict["elapsed"],
)
self.step_wise_images.append(stepwise_img)
stepwise_img.save(
path=self.output_dir / f"seed_{seed}_step{step}of{len(time_steps)}.png",
export_json_metadata=False,
)
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)
composite_img.save(self.output_dir / f"seed_{seed}_composite.png")

23
src/mflux/controlnet/controlnet_util.py
View File

@ -1,12 +1,10 @@
import logging
import os
import cv2
import mlx.core as mx
import numpy as np
import PIL.Image
from mflux.config.runtime_config import RuntimeConfig
from mflux.models.vae.vae import VAE
from mflux.post_processing.array_util import ArrayUtil
@ -17,29 +15,20 @@ class ControlnetUtil:
@staticmethod
def encode_image(
vae: VAE,
config: RuntimeConfig,
height: int,
width: int,
controlnet_image_path: str,
controlnet_save_canny: bool,
output: str,
) -> mx.array:
) -> (mx.array, PIL.Image):
from mflux import ImageUtil
control_image = ImageUtil.load_image(controlnet_image_path)
control_image = ControlnetUtil._scale_image(config.height, config.width, control_image)
control_image = ControlnetUtil._scale_image(height=height, width=width, img=control_image)
control_image = ControlnetUtil._preprocess_canny(control_image)
if controlnet_save_canny:
base, ext = os.path.splitext(output)
ImageUtil.save_image(
image=control_image,
path=f"{base}_controlnet_canny{ext}"
) # fmt: off
controlnet_cond = ImageUtil.to_array(control_image)
controlnet_cond = vae.encode(controlnet_cond)
controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor
controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=config.height, width=config.width)
return controlnet_cond
controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=height, width=width)
return controlnet_cond, control_image
@staticmethod
def _preprocess_canny(img: PIL.Image) -> PIL.Image:

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

@ -1,15 +1,13 @@
from pathlib import Path
import mlx.core as mx
from mlx import nn
from tqdm import tqdm
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.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
@ -19,7 +17,6 @@ from mflux.models.vae.vae import VAE
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.generated_image import GeneratedImage
from mflux.post_processing.image_util import ImageUtil
from mflux.post_processing.stepwise_handler import StepwiseHandler
from mflux.weights.model_saver import ModelSaver
@ -53,49 +50,44 @@ class Flux1Controlnet(nn.Module):
self,
seed: int,
prompt: str,
output: str,
controlnet_image_path: str,
controlnet_save_canny: bool = False,
config: Config = Config(),
stepwise_output_dir: Path = None,
) -> GeneratedImage:
# Convert the user config to a runtime config with derived parameters.
# 0. Create a new runtime config based on the model type and input parameters
config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.num_inference_steps))
stepwise_handler = StepwiseHandler(
flux=self,
config=config,
seed=seed,
prompt=prompt,
time_steps=time_steps,
output_dir=stepwise_output_dir,
)
# 0. Encode the controlnet reference image
controlnet_condition = ControlnetUtil.encode_image(
# 1. Encode the controlnet reference image
controlnet_condition, canny_image = ControlnetUtil.encode_image(
vae=self.vae,
config=config,
height=config.height,
width=config.width,
controlnet_image_path=controlnet_image_path,
controlnet_save_canny=controlnet_save_canny,
output=output,
)
# 1. Create the initial latents
# 2. Create the initial latents
latents = LatentCreator.create(
seed=seed,
height=config.height,
width=config.width
) # fmt: off
# 2. Embed the prompt
# 3. Encode the prompt
t5_tokens = self.t5_tokenizer.tokenize(prompt)
clip_tokens = self.clip_tokenizer.tokenize(prompt)
prompt_embeds = self.t5_text_encoder(t5_tokens)
pooled_prompt_embeds = self.clip_text_encoder(clip_tokens)
# (Optional) Call subscribers for beginning of loop
Callbacks.before_loop(
seed=seed,
prompt=prompt,
canny_image=canny_image
) # fmt: off
for gen_step, t in enumerate(time_steps, 1):
try:
# 3.t Compute controlnet samples
# 4.t Compute controlnet samples
controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet(
t=t,
config=config,
@ -105,7 +97,7 @@ class Flux1Controlnet(nn.Module):
controlnet_condition=controlnet_condition,
)
# 4.t Predict the noise
# 5.t Predict the noise
noise = self.transformer(
t=t,
config=config,
@ -116,21 +108,34 @@ class Flux1Controlnet(nn.Module):
controlnet_single_block_samples=controlnet_single_block_samples,
)
# 5.t Take one denoise step
# 6.t Take one denoise step
dt = config.sigmas[t + 1] - config.sigmas[t]
latents += noise * dt
# Handle stepwise output if enabled
stepwise_handler.process_step(gen_step, latents)
# (Optional) Call subscribes at end of loop
Callbacks.in_loop(
seed=seed,
prompt=prompt,
step=gen_step,
latents=latents,
config=config,
time_steps=time_steps,
) # fmt: off
# Evaluate to enable progress tracking
# (Optional) Evaluate to enable progress tracking
mx.eval(latents)
except KeyboardInterrupt: # noqa: PERF203
stepwise_handler.handle_interruption()
raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
Callbacks.interruption(
seed=seed,
prompt=prompt,
step=gen_step,
latents=latents,
config=config,
time_steps=time_steps,
)
# 5. Decode the latent array and return the image
# 7. 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)
return ImageUtil.to_image(

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

@ -1,16 +1,13 @@
import warnings
from pathlib import Path
import mlx.core as mx
from mlx import nn
from tqdm import tqdm
from mflux.callbacks.callbacks import Callbacks
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig, ModelLookup
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 LatentCreator
from mflux.latent_creator.latent_creator import Img2Img, LatentCreator
from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
from mflux.models.transformer.transformer import Transformer
@ -18,7 +15,6 @@ from mflux.models.vae.vae import VAE
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.generated_image import GeneratedImage
from mflux.post_processing.image_util import ImageUtil
from mflux.post_processing.stepwise_handler import StepwiseHandler
from mflux.weights.model_saver import ModelSaver
@ -51,33 +47,36 @@ class Flux1(nn.Module):
seed: int,
prompt: str,
config: Config = Config(),
stepwise_output_dir: Path = None,
) -> GeneratedImage:
# Convert the user config to a runtime config with derived parameters.
# 0. Create a new runtime config based on the model type and input parameters
config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
stepwise_handler = StepwiseHandler(
flux=self,
config=config,
seed=seed,
prompt=prompt,
time_steps=time_steps,
output_dir=stepwise_output_dir,
)
# 1. Create the initial latents
latents = LatentCreator.create_for_txt2img_or_img2img(
seed=seed,
vae=self.vae,
runtime_conf=config,
height=config.height,
width=config.width,
img2img=Img2Img(
vae=self.vae,
sigmas=config.sigmas,
init_time_step=config.init_time_step,
init_image_path=config.init_image_path,
),
)
# 2. Embed the prompt
# 2. Encode the prompt
t5_tokens = self.t5_tokenizer.tokenize(prompt)
clip_tokens = self.clip_tokenizer.tokenize(prompt)
prompt_embeds = self.t5_text_encoder(t5_tokens)
pooled_prompt_embeds = self.clip_text_encoder(clip_tokens)
# (Optional) Call subscribers for beginning of loop
Callbacks.before_loop(
seed=seed,
prompt=prompt
) # fmt: off
for gen_step, t in enumerate(time_steps, 1):
try:
# 3.t Predict the noise
@ -93,17 +92,30 @@ class Flux1(nn.Module):
dt = config.sigmas[t + 1] - config.sigmas[t]
latents += noise * dt
# Handle stepwise output if enabled
stepwise_handler.process_step(gen_step, latents)
# (Optional) Call subscribes at end of loop
Callbacks.in_loop(
seed=seed,
prompt=prompt,
step=gen_step,
latents=latents,
config=config,
time_steps=time_steps,
) # fmt: off
# Evaluate to enable progress tracking
# (Optional) Evaluate to enable progress tracking
mx.eval(latents)
except KeyboardInterrupt: # noqa: PERF203
stepwise_handler.handle_interruption()
raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
Callbacks.interruption(
seed=seed,
prompt=prompt,
step=gen_step,
latents=latents,
config=config,
time_steps=time_steps,
)
# 5. Decode the latent array and return the image
# 7. 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)
return ImageUtil.to_image(
@ -119,15 +131,6 @@ class Flux1(nn.Module):
generation_time=time_steps.format_dict["elapsed"],
)
@staticmethod
def from_alias(alias: str, quantize: int | None = None) -> "Flux1":
warnings.warn(
"from_alias is deprecated and will be removed in a future release. Please use from_name instead.",
DeprecationWarning,
stacklevel=2,
)
return Flux1.from_name(model_name=alias, quantize=quantize)
@staticmethod
def from_name(model_name: str, quantize: int | None = None) -> "Flux1":
return Flux1(

23
src/mflux/generate.py
View File

@ -1,11 +1,11 @@
from pathlib import Path
from mflux import Config, Flux1, ModelLookup, StopImageGenerationException
from mflux.callbacks.callback_registry import CallbackRegistry
from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
from mflux.ui.cli.parsers import CommandLineParser
def main():
# fmt: off
# 0. Parse command line arguments
parser = CommandLineParser(description="Generate an image based on a prompt.")
parser.add_model_arguments(require_model_arg=False)
parser.add_lora_arguments()
@ -23,13 +23,18 @@ def main():
lora_scales=args.lora_scales,
)
# 2. Register the optional callbacks
if args.stepwise_image_output_dir:
handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
CallbackRegistry.register_in_loop(handler)
CallbackRegistry.register_interrupt(handler)
try:
for seed_value in args.seed:
# 2. Generate an image for each seed value
for seed in args.seed:
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed_value,
seed=seed,
prompt=args.prompt,
stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
config=Config(
num_inference_steps=args.steps,
height=args.height,
@ -39,8 +44,8 @@ def main():
init_image_strength=args.init_image_strength,
),
)
# 3. Save the image
image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata)
# 4. Save the image
image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
except StopImageGenerationException as stop_exc:
print(stop_exc)

27
src/mflux/generate_controlnet.py
View File

@ -1,10 +1,12 @@
from pathlib import Path
from mflux import Config, Flux1Controlnet, ModelLookup, StopImageGenerationException
from mflux.callbacks.callback_registry import CallbackRegistry
from mflux.callbacks.instances.canny_saver import CannyImageSaver
from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
from mflux.ui.cli.parsers import CommandLineParser
def main():
# 0. Parse command line arguments
parser = CommandLineParser(description="Generate an image based on a prompt and a controlnet reference image.") # fmt: off
parser.add_model_arguments(require_model_arg=True)
parser.add_lora_arguments()
@ -22,16 +24,21 @@ def main():
lora_scales=args.lora_scales,
)
# 2. Register the optional callbacks
if args.controlnet_save_canny:
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_in_loop(handler)
CallbackRegistry.register_interrupt(handler)
try:
for seed_value in args.seed:
# 2. Generate an image for each seed value
for seed in args.seed:
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed_value,
seed=seed,
prompt=args.prompt,
output=args.output,
controlnet_image_path=args.controlnet_image_path,
controlnet_save_canny=args.controlnet_save_canny,
stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
config=Config(
num_inference_steps=args.steps,
height=args.height,
@ -41,8 +48,8 @@ def main():
),
)
# 3. Save the image
image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata)
# 4. Save the image
image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
except StopImageGenerationException as stop_exc:
print(stop_exc)

69
src/mflux/latent_creator/latent_creator.py
View File

@ -1,11 +1,24 @@
import mlx.core as mx
from mlx import nn
from mflux.config.runtime_config import RuntimeConfig
from mflux.models.vae.vae import VAE
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.image_util import ImageUtil
class Img2Img:
def __init__(
self,
vae: VAE,
sigmas: mx.array,
init_time_step: int,
init_image_path: int,
):
self.vae = vae
self.sigmas = sigmas
self.init_time_step = init_time_step
self.init_image_path = init_image_path
class LatentCreator:
@staticmethod
def create(
@ -21,29 +34,41 @@ class LatentCreator:
@staticmethod
def create_for_txt2img_or_img2img(
seed: int,
runtime_conf: RuntimeConfig,
vae: nn.Module,
height: int,
width: int,
img2img: Img2Img,
) -> mx.array:
pure_noise = LatentCreator.create(
seed=seed,
height=runtime_conf.height,
width=runtime_conf.width,
)
# 0. Determine type of image generation
is_text2img = img2img.init_image_path is None
if runtime_conf.config.init_image_path is None:
# Text2Image
return pure_noise
else:
# Image2Image
user_image = ImageUtil.load_image(runtime_conf.config.init_image_path).convert("RGB")
scaled_user_image = ImageUtil.scale_to_dimensions(
image=user_image,
target_width=runtime_conf.width,
target_height=runtime_conf.height,
if is_text2img:
# 1. Create the pure noise
return LatentCreator.create(
seed=seed,
height=height,
width=width,
)
encoded = vae.encode(ImageUtil.to_array(scaled_user_image))
latents = ArrayUtil.pack_latents(latents=encoded, height=runtime_conf.height, width=runtime_conf.width)
sigma = runtime_conf.sigmas[runtime_conf.init_time_step]
else:
# 1. Create the pure noise
pure_noise = LatentCreator.create(
seed=seed,
height=height,
width=width,
)
# 2. Encode the image
scaled_user_image = ImageUtil.scale_to_dimensions(
image=ImageUtil.load_image(img2img.init_image_path).convert("RGB"),
target_width=width,
target_height=height,
)
encoded = img2img.vae.encode(ImageUtil.to_array(scaled_user_image))
latents = ArrayUtil.pack_latents(latents=encoded, height=height, width=width)
# 3. Find the appropriate sigma value
sigma = img2img.sigmas[img2img.init_time_step]
# 4. Blend the appropriate amount of noise based on linear interpolation
return LatentCreator.add_noise_by_interpolation(
clean=latents,
noise=pure_noise,

60
src/mflux/post_processing/stepwise_handler.py
View File

@ -1,60 +0,0 @@
from pathlib import Path
import mlx.core as mx
import tqdm
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:
def __init__(
self,
flux,
config: RuntimeConfig,
seed: int,
prompt: str,
time_steps: tqdm.std.tqdm,
output_dir: Path | None = None,
):
self.flux = flux
self.config = config
self.seed = seed
self.prompt = prompt
self.output_dir = output_dir
self.time_steps = time_steps
self.step_wise_images = []
if self.output_dir:
self.output_dir.mkdir(parents=True, exist_ok=True)
def save_composite(self):
if self.step_wise_images:
composite_img = ImageUtil.to_composite_image(self.step_wise_images)
composite_img.save(self.output_dir / f"seed_{self.seed}_composite.png")
def process_step(self, gen_step: int, latents: mx.array):
if self.output_dir:
unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=self.config.height, width=self.config.width) # fmt: off
stepwise_decoded = self.flux.vae.decode(unpack_latents)
stepwise_img = ImageUtil.to_image(
decoded_latents=stepwise_decoded,
config=self.config,
seed=self.seed,
prompt=self.prompt,
quantization=self.flux.bits,
lora_paths=self.flux.lora_paths,
lora_scales=self.flux.lora_scales,
generation_time=self.time_steps.format_dict["elapsed"],
)
self.step_wise_images.append(stepwise_img)
stepwise_img.save(
path=self.output_dir / f"seed_{self.seed}_step{gen_step}of{len(self.time_steps)}.png",
export_json_metadata=False,
)
self.save_composite()
def handle_interruption(self):
self.save_composite()

2
tests/image_generation/helpers/image_generation_controlnet_test_helper.py
View File

@ -40,9 +40,7 @@ class ImageGeneratorControlnetTestHelper:
image = flux.generate_image(
seed=seed,
prompt=prompt,
output=str(output_image_path),
controlnet_image_path=controlnet_image_path,
controlnet_save_canny=False,
config=Config(
num_inference_steps=steps,
height=768,