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

Merge pull request #124 from filipstrand/refactor-and-add-callback-mechanism

Browse Source 
Refactor and add callback mechanism
This commit is contained in:
Filip Strand 2025-02-10 18:53:52 +01:00 committed by GitHub
commit 2d481e77b2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
20 changed files with 585 additions and 323 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/mflux/__init__.py
View File

@ -1,4 +1,4 @@
from mflux.config.config import Config, ConfigControlnet
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig, ModelLookup
from mflux.controlnet.flux_controlnet import Flux1Controlnet
from mflux.error.exceptions import StopImageGenerationException
@ -9,7 +9,6 @@ __all__ = [
"Flux1",
"Flux1Controlnet",
"Config",
"ConfigControlnet",
"ModelConfig",
"ModelLookup",
"ImageUtil",

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")

13
src/mflux/config/config.py
View File

@ -17,6 +17,7 @@ class Config:
guidance: float = 4.0,
init_image_path: Path | None = None,
init_image_strength: float | None = None,
controlnet_strength: float | None = None,
):
if width % 16 != 0 or height % 16 != 0:
log.warning("Width and height should be multiples of 16. Rounding down.")
@ -26,16 +27,4 @@ class Config:
self.guidance = guidance
self.init_image_path = init_image_path
self.init_image_strength = init_image_strength
class ConfigControlnet(Config):
def __init__(
self,
num_inference_steps: int = 4,
width: int = 1024,
height: int = 1024,
guidance: float = 4.0,
controlnet_strength: float = 1.0,
):
super().__init__(num_inference_steps=num_inference_steps, width=width, height=height, guidance=guidance)
self.controlnet_strength = controlnet_strength

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

@ -3,14 +3,18 @@ import logging
import mlx.core as mx
import numpy as np
from mflux.config.config import Config, ConfigControlnet
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
logger = logging.getLogger(__name__)
class RuntimeConfig:
def __init__(self, config: Config | ConfigControlnet, model_config):
def __init__(
self,
config: Config,
model_config: ModelConfig,
):
self.config = config
self.model_config = model_config
self.sigmas = self._create_sigmas(config, model_config)
@ -50,23 +54,20 @@ class RuntimeConfig:
@property
def init_time_step(self) -> int:
if self.config.init_image_path is None:
# text to image, always begin at time step 0
# For text-to-image, always begin at time step 0.
return 0
else:
# we skip to the time step as informed by the init_image_strength
# the higher the strength number, the more time steps we skip
# For image-to-image: higher strength means we skip more steps.
strength = max(0.0, min(1.0, self.config.init_image_strength))
# if the strength is too small to even influence the image
# help the user round up so the init_image has influence at step 1
t = max(1, int(self.num_inference_steps * strength))
return t
@property
def controlnet_strength(self) -> float:
if isinstance(self.config, ConfigControlnet):
def controlnet_strength(self) -> float | None:
if self.config.controlnet_strength is not None:
return self.config.controlnet_strength
else:
raise NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet")
return None
@staticmethod
def _create_sigmas(config, model) -> mx.array:

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

@ -1,16 +1,37 @@
import logging
import os
import cv2
import mlx.core as mx
import numpy as np
import PIL.Image
from mflux.models.vae.vae import VAE
from mflux.post_processing.array_util import ArrayUtil
log = logging.getLogger(__name__)
class ControlnetUtil:
@staticmethod
def preprocess_canny(img: PIL.Image) -> PIL.Image:
def encode_image(
vae: VAE,
height: int,
width: int,
controlnet_image_path: str,
) -> (mx.array, PIL.Image):
from mflux import ImageUtil
control_image = ImageUtil.load_image(controlnet_image_path)
control_image = ControlnetUtil._scale_image(height=height, width=width, img=control_image)
control_image = ControlnetUtil._preprocess_canny(control_image)
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=height, width=width)
return controlnet_cond, control_image
@staticmethod
def _preprocess_canny(img: PIL.Image) -> PIL.Image:
image_to_canny = np.array(img)
image_to_canny = cv2.Canny(image_to_canny, 100, 200)
image_to_canny = np.array(image_to_canny[:, :, None])
@ -18,16 +39,8 @@ class ControlnetUtil:
return PIL.Image.fromarray(image_to_canny)
@staticmethod
def scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
def _scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
if height != img.height or width != img.width:
log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}")
img = img.resize((width, height), PIL.Image.LANCZOS)
return img
@staticmethod
def save_canny_image(control_image: PIL.Image, path: str):
from mflux import ImageUtil
base, ext = os.path.splitext(path)
new_filename = f"{base}_controlnet_canny{ext}"
ImageUtil.save_image(control_image, new_filename)

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

@ -1,15 +1,14 @@
from pathlib import Path
import mlx.core as mx
from mlx import nn
from tqdm import tqdm
from mflux.config.config import ConfigControlnet
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.controlnet.weight_handler_controlnet import WeightHandlerControlnet
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
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
@ -18,17 +17,16 @@ 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.tokenizer.clip_tokenizer import TokenizerCLIP
from mflux.tokenizer.t5_tokenizer import TokenizerT5
from mflux.tokenizer.tokenizer_handler import TokenizerHandler
from mflux.weights.model_saver import ModelSaver
from mflux.weights.weight_handler import WeightHandler
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
from mflux.weights.weight_util import WeightUtil
class Flux1Controlnet:
class Flux1Controlnet(nn.Module):
vae: VAE
transformer: Transformer
transformer_controlnet: TransformerControlnet
t5_text_encoder: T5Encoder
clip_text_encoder: CLIPEncoder
def __init__(
self,
model_config: ModelConfig,
@ -38,84 +36,58 @@ class Flux1Controlnet:
lora_scales: list[float] | None = None,
controlnet_path: str | None = None,
):
self.lora_paths = lora_paths
self.lora_scales = lora_scales
self.model_config = model_config
# Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
# Load the weights
weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path)
# Initialize the models
self.vae = VAE()
self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks()) # fmt: off
self.t5_text_encoder = T5Encoder()
self.clip_text_encoder = CLIPEncoder()
# Set the weights and quantize the model
self.bits = WeightUtil.set_weights_and_quantize(
quantize_arg=quantize,
weights=weights,
vae=self.vae,
transformer=self.transformer,
t5_text_encoder=self.t5_text_encoder,
clip_text_encoder=self.clip_text_encoder,
)
# Set LoRA weights
lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales) # fmt:off
WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
# Set Controlnet weights
weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
self.transformer_controlnet = TransformerControlnet(model_config=model_config, num_transformer_blocks=weights_controlnet.num_transformer_blocks(), num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks()) # fmt:off
WeightUtil.set_controlnet_weights_and_quantize(
quantize_arg=quantize,
weights=weights_controlnet,
transformer_controlnet=self.transformer_controlnet,
super().__init__()
FluxInitializer.init_controlnet(
flux_model=self,
model_config=model_config,
quantize=quantize,
local_path=local_path,
lora_paths=lora_paths,
lora_scales=lora_scales,
)
def generate_image(
self,
seed: int,
prompt: str,
output: str,
controlnet_image_path: str,
controlnet_save_canny: bool = False,
config: ConfigControlnet = ConfigControlnet(),
stepwise_output_dir: Path = None,
) -> GeneratedImage: # fmt: off
# Create a new runtime config based on the model type and input parameters
config: Config = Config(),
) -> GeneratedImage:
# 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,
# 1. Encode the controlnet reference image
controlnet_condition, canny_image = ControlnetUtil.encode_image(
vae=self.vae,
height=config.height,
width=config.width,
controlnet_image_path=controlnet_image_path,
)
# 0. Embed the controlnet reference image
controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output)
# 2. Create the initial latents
latents = LatentCreator.create(
seed=seed,
height=config.height,
width=config.width
) # fmt: off
# 1. Create the initial latents
latents = LatentCreator.create(seed=seed, height=config.height, width=config.width)
# 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,
@ -125,7 +97,7 @@ class Flux1Controlnet:
controlnet_condition=controlnet_condition,
)
# 4.t Predict the noise
# 5.t Predict the noise
noise = self.transformer(
t=t,
config=config,
@ -136,55 +108,48 @@ class Flux1Controlnet:
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(
decoded_latents=decoded,
config=config,
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict["elapsed"],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
config=config,
controlnet_image_path=controlnet_image_path,
generation_time=time_steps.format_dict["elapsed"],
)
def _embed_image(
self,
config: RuntimeConfig,
controlnet_image_path: str,
controlnet_save_canny: bool,
output: str,
):
control_image = ImageUtil.load_image(controlnet_image_path)
control_image = ControlnetUtil.scale_image(config.height, config.width, control_image)
control_image = ControlnetUtil.preprocess_canny(control_image)
if controlnet_save_canny:
ControlnetUtil.save_canny_image(control_image, output)
controlnet_cond = ImageUtil.to_array(control_image)
controlnet_cond = self.vae.encode(controlnet_cond)
controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor
controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=config.height, width=config.width)
return controlnet_cond
def save_model(self, base_path: str) -> None:
ModelSaver.save_model(self, self.bits, base_path)
ModelSaver.save_weights(base_path, self.bits, self.transformer_controlnet, "transformer_controlnet")

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

@ -1,15 +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.latent_creator.latent_creator import LatentCreator
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
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
from mflux.models.transformer.transformer import Transformer
@ -17,17 +15,15 @@ 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.tokenizer.clip_tokenizer import TokenizerCLIP
from mflux.tokenizer.t5_tokenizer import TokenizerT5
from mflux.tokenizer.tokenizer_handler import TokenizerHandler
from mflux.weights.model_saver import ModelSaver
from mflux.weights.weight_handler import WeightHandler
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
from mflux.weights.weight_util import WeightUtil
class Flux1(nn.Module):
vae: VAE
transformer: Transformer
t5_text_encoder: T5Encoder
clip_text_encoder: CLIPEncoder
def __init__(
self,
model_config: ModelConfig,
@ -37,66 +33,50 @@ class Flux1(nn.Module):
lora_scales: list[float] | None = None,
):
super().__init__()
self.lora_paths = lora_paths
self.lora_scales = lora_scales
self.model_config = model_config
# Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
# Load the weights
weights = WeightHandler.load_regular_weights(repo_id=model_config.model_name, local_path=local_path)
# Initialize the models
self.vae = VAE()
self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks()) # fmt: off
self.t5_text_encoder = T5Encoder()
self.clip_text_encoder = CLIPEncoder()
# Set the weights and quantize the model
self.bits = WeightUtil.set_weights_and_quantize(
quantize_arg=quantize,
weights=weights,
vae=self.vae,
transformer=self.transformer,
t5_text_encoder=self.t5_text_encoder,
clip_text_encoder=self.clip_text_encoder,
FluxInitializer.init(
flux_model=self,
model_config=model_config,
quantize=quantize,
local_path=local_path,
lora_paths=lora_paths,
lora_scales=lora_scales,
)
# Set LoRA weights
lora_weights = WeightHandlerLoRA.load_lora_weights(transformer=self.transformer, lora_files=lora_paths, lora_scales=lora_scales) # fmt:off
WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
def generate_image(
self,
seed: int,
prompt: str,
config: Config = Config(),
stepwise_output_dir: Path = None,
) -> GeneratedImage:
# Create a new runtime config based on the model type and input 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, config, self.vae)
latents = LatentCreator.create_for_txt2img_or_img2img(
seed=seed,
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
@ -112,41 +92,45 @@ 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(
decoded_latents=decoded,
config=config,
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict["elapsed"],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
init_image_path=config.init_image_path,
init_image_strength=config.init_image_strength,
config=config,
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(

111
src/mflux/flux/flux_initializer.py Normal file
View File

@ -0,0 +1,111 @@
from mflux.controlnet.transformer_controlnet import TransformerControlnet
from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
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
from mflux.models.vae.vae import VAE
from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
from mflux.tokenizer.t5_tokenizer import TokenizerT5
from mflux.tokenizer.tokenizer_handler import TokenizerHandler
from mflux.weights.weight_handler import WeightHandler
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
from mflux.weights.weight_util import WeightUtil
class FluxInitializer:
@staticmethod
def init(
flux_model,
model_config,
quantize: int | None,
local_path: str | None,
lora_paths: list[str] | None,
lora_scales: list[float] | None,
) -> None:
# 0. Set paths and config for later
flux_model.lora_paths = lora_paths
flux_model.lora_scales = lora_scales
flux_model.model_config = model_config
# 1. Initialize tokenizers
tokenizers = TokenizerHandler(
repo_id=model_config.model_name,
max_t5_length=model_config.max_sequence_length,
local_path=local_path,
)
flux_model.t5_tokenizer = TokenizerT5(
tokenizer=tokenizers.t5,
max_length=model_config.max_sequence_length
) # fmt: off
flux_model.clip_tokenizer = TokenizerCLIP(
tokenizer=tokenizers.clip,
)
# 2. Load the regular weights
weights = WeightHandler.load_regular_weights(
repo_id=model_config.model_name,
local_path=local_path
) # fmt: off
# 3. Initialize all models
flux_model.vae = VAE()
flux_model.transformer = Transformer(
model_config=model_config,
num_transformer_blocks=weights.num_transformer_blocks(),
num_single_transformer_blocks=weights.num_single_transformer_blocks(),
)
flux_model.t5_text_encoder = T5Encoder()
flux_model.clip_text_encoder = CLIPEncoder()
# 4. Apply weights and quantize the models
flux_model.bits = WeightUtil.set_weights_and_quantize(
quantize_arg=quantize,
weights=weights,
vae=flux_model.vae,
transformer=flux_model.transformer,
t5_text_encoder=flux_model.t5_text_encoder,
clip_text_encoder=flux_model.clip_text_encoder,
)
# 5. Set LoRA weights
lora_weights = WeightHandlerLoRA.load_lora_weights(
transformer=flux_model.transformer,
lora_files=lora_paths,
lora_scales=lora_scales,
)
WeightHandlerLoRA.set_lora_weights(
transformer=flux_model.transformer,
loras=lora_weights
) # fmt: off
@staticmethod
def init_controlnet(
flux_model,
model_config,
quantize: int | None,
local_path: str | None,
lora_paths: list[str] | None,
lora_scales: list[float] | None,
) -> None:
# 1. Start with same init as regular Flux
FluxInitializer.init(
flux_model=flux_model,
model_config=model_config,
quantize=quantize,
local_path=local_path,
lora_paths=lora_paths,
lora_scales=lora_scales,
)
# 2. Apply ControlNet-specific initialization
weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
flux_model.transformer_controlnet = TransformerControlnet(
model_config=model_config,
num_transformer_blocks=weights_controlnet.num_transformer_blocks(),
num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks(),
)
WeightUtil.set_controlnet_weights_and_quantize(
quantize_arg=quantize,
weights=weights_controlnet,
transformer_controlnet=flux_model.transformer_controlnet,
)

25
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()
@ -14,7 +14,7 @@ def main():
parser.add_output_arguments()
args = parser.parse_args()
# Load the model
# 1. Load the model
flux = Flux1(
model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model),
quantize=args.quantize,
@ -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:
# 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,
),
)
# 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)

33
src/mflux/generate_controlnet.py
View File

@ -1,10 +1,12 @@
from pathlib import Path
from mflux import ConfigControlnet, Flux1Controlnet, ModelLookup, StopImageGenerationException
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()
@ -13,7 +15,7 @@ def main():
parser.add_output_arguments()
args = parser.parse_args()
# Load the model
# 1. Load the model
flux = Flux1Controlnet(
model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model),
quantize=args.quantize,
@ -22,17 +24,22 @@ 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:
# 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=ConfigControlnet(
config=Config(
num_inference_steps=args.steps,
height=args.height,
width=args.width,
@ -41,8 +48,8 @@ def main():
),
)
# 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,

8
src/mflux/post_processing/image_util.py
View File

@ -8,25 +8,23 @@ import numpy as np
import piexif
import PIL.Image
from mflux.config.config import ConfigControlnet
from mflux.config.runtime_config import RuntimeConfig
from mflux.post_processing.generated_image import GeneratedImage
log = logging.getLogger(__name__)
RuntimeConfig: t.TypeAlias = "mflux.config.runtime_config.RuntimeConfig" # noqa: F821
class ImageUtil:
@staticmethod
def to_image(
decoded_latents: mx.array,
config: RuntimeConfig,
seed: int,
prompt: str,
quantization: int,
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
config: RuntimeConfig,
controlnet_image_path: str | None = None,
init_image_path: str | None = None,
init_image_strength: float | None = None,
@ -49,7 +47,7 @@ class ImageUtil:
init_image_path=init_image_path,
init_image_strength=init_image_strength,
controlnet_image_path=controlnet_image_path,
controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None,
controlnet_strength=config.controlnet_strength,
)
@staticmethod

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,
seed=self.seed,
prompt=self.prompt,
quantization=self.flux.bits,
generation_time=self.time_steps.format_dict["elapsed"],
lora_paths=self.flux.lora_paths,
lora_scales=self.flux.lora_scales,
config=self.config,
)
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()

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

@ -3,7 +3,7 @@ import os
import numpy as np
from PIL import Image
from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig
from mflux import Config, Flux1Controlnet, ModelConfig
from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
@ -40,10 +40,8 @@ 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=ConfigControlnet(
config=Config(
num_inference_steps=steps,
height=768,
width=493,