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

Refactor: Add separate Flux initializer

Browse Source
This commit is contained in:
filipstrand 2025-02-08 05:09:33 +01:00
parent 8f82549f43
commit 22f98fc5a4
5 changed files with 199 additions and 119 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -10,7 +10,11 @@ logger = logging.getLogger(__name__)
class RuntimeConfig:
def __init__(self, config: Config | ConfigControlnet, model_config):
def __init__(
self,
config: Config | ConfigControlnet,
model_config: ModelConfig,
):
self.config = config
self.model_config = model_config
self.sigmas = self._create_sigmas(config, model_config)

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

@ -2,15 +2,47 @@ 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
log = logging.getLogger(__name__)
class ControlnetUtil:
@staticmethod
def preprocess_canny(img: PIL.Image) -> PIL.Image:
def encode_image(
vae: VAE,
config: RuntimeConfig,
controlnet_image_path: str,
controlnet_save_canny: bool,
output: str,
) -> mx.array:
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._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
@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 +50,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)

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

@ -1,6 +1,7 @@
from pathlib import Path
import mlx.core as mx
from mlx import nn
from tqdm import tqdm
from mflux.config.config import ConfigControlnet
@ -8,8 +9,8 @@ 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
@ -19,16 +20,16 @@ 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,45 +39,14 @@ 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(
@ -88,7 +58,7 @@ class Flux1Controlnet:
controlnet_save_canny: bool = False,
config: ConfigControlnet = ConfigControlnet(),
stepwise_output_dir: Path = None,
) -> GeneratedImage: # fmt: off
) -> GeneratedImage:
# 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))
@ -101,11 +71,21 @@ class Flux1Controlnet:
output_dir=stepwise_output_dir,
)
# 0. Embed the controlnet reference image
controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output)
# 0. Encode the controlnet reference image
controlnet_condition = ControlnetUtil.encode_image(
vae=self.vae,
config=config,
controlnet_image_path=controlnet_image_path,
controlnet_save_canny=controlnet_save_canny,
output=output,
)
# 1. Create the initial latents
latents = LatentCreator.create(seed=seed, height=config.height, width=config.width)
latents = LatentCreator.create(
seed=seed,
height=config.height,
width=config.width
) # fmt: off
# 2. Embed the prompt
t5_tokens = self.t5_tokenizer.tokenize(prompt)
@ -165,26 +145,6 @@ class Flux1Controlnet:
controlnet_image_path=controlnet_image_path,
)
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")

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

@ -9,6 +9,7 @@ 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.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
@ -18,16 +19,15 @@ 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,38 +37,15 @@ 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,
@ -89,7 +66,11 @@ class Flux1(nn.Module):
)
# 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,
vae=self.vae,
runtime_conf=config,
)
# 2. Embed the prompt
t5_tokens = self.t5_tokenizer.tokenize(prompt)

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