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

Refactor: Controlnet transformer & Flux controlnet

Browse Source
This commit is contained in:
filipstrand 2025-02-01 21:22:47 +01:00
parent 010610e33b
commit fa99cf0cc8
3 changed files with 144 additions and 100 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,6 +1,4 @@
import logging
from pathlib import Path
from typing import TYPE_CHECKING
import mlx.core as mx
from tqdm import tqdm
@ -18,6 +16,7 @@ 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.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
@ -28,14 +27,6 @@ from mflux.weights.weight_handler import WeightHandler
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
from mflux.weights.weight_util import WeightUtil
if TYPE_CHECKING:
from mflux.post_processing.generated_image import GeneratedImage
log = logging.getLogger(__name__)
CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
class Flux1Controlnet:
def __init__(
@ -61,7 +52,7 @@ class Flux1Controlnet:
# Initialize the models
self.vae = VAE()
self.transformer = Transformer(model_config, num_transformer_blocks=weights.num_transformer_blocks())
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()
@ -80,8 +71,8 @@ class Flux1Controlnet:
WeightHandlerLoRA.set_lora_weights(transformer=self.transformer, loras=lora_weights)
# Set Controlnet weights
weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer(controlnet_id=CONTROLNET_ID)
self.transformer_controlnet = TransformerControlnet(model_config=model_config, num_blocks=weights_controlnet.config["num_layers"], num_single_blocks=weights_controlnet.config["num_single_layers"]) # fmt:off
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,
@ -89,15 +80,15 @@ class Flux1Controlnet:
)
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
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 = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.num_inference_steps))
@ -110,16 +101,8 @@ class Flux1Controlnet:
output_dir=stepwise_output_dir,
)
# Embed the controlnet reference image
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)
# 0. Embed the controlnet reference image
controlnet_condition = self._embed_image(config, controlnet_image_path, controlnet_save_canny, output)
# 1. Create the initial latents
latents = LatentCreator.create(seed=seed, height=config.height, width=config.width)
@ -130,35 +113,35 @@ class Flux1Controlnet:
prompt_embeds = self.t5_text_encoder(t5_tokens)
pooled_prompt_embeds = self.clip_text_encoder(clip_tokens)
for t in time_steps:
for gen_step, t in enumerate(time_steps, 1):
try:
# Compute controlnet samples
# 3.t Compute controlnet samples
controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet(
t=t,
config=config,
hidden_states=latents,
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
hidden_states=latents,
controlnet_cond=controlnet_cond,
config=config,
controlnet_condition=controlnet_condition,
)
# 3.t Predict the noise
noise = self.transformer.predict(
# 4.t Predict the noise
noise = self.transformer(
t=t,
config=config,
hidden_states=latents,
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
hidden_states=latents,
config=config,
controlnet_block_samples=controlnet_block_samples,
controlnet_single_block_samples=controlnet_single_block_samples,
)
# 4.t Take one denoise step
# 5.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(t, latents)
stepwise_handler.process_step(gen_step, latents)
# Evaluate to enable progress tracking
mx.eval(latents)
@ -182,6 +165,26 @@ 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")

143
src/mflux/controlnet/transformer_controlnet.py
View File

@ -4,12 +4,8 @@ from mlx import nn
from mflux.config.model_config import ModelConfig
from mflux.config.runtime_config import RuntimeConfig
from mflux.models.transformer.embed_nd import EmbedND
from mflux.models.transformer.joint_transformer_block import (
JointTransformerBlock,
)
from mflux.models.transformer.single_transformer_block import (
SingleTransformerBlock,
)
from mflux.models.transformer.joint_transformer_block import JointTransformerBlock
from mflux.models.transformer.single_transformer_block import SingleTransformerBlock
from mflux.models.transformer.time_text_embed import TimeTextEmbed
from mflux.models.transformer.transformer import Transformer
@ -18,80 +14,117 @@ class TransformerControlnet(nn.Module):
def __init__(
self,
model_config: ModelConfig,
num_blocks: int,
num_single_blocks: int,
num_transformer_blocks: int = 5,
num_single_transformer_blocks: int = 0,
):
super().__init__()
self.pos_embed = EmbedND()
self.x_embedder = nn.Linear(64, 3072)
self.time_text_embed = TimeTextEmbed(model_config=model_config)
self.context_embedder = nn.Linear(4096, 3072)
self.transformer_blocks = [JointTransformerBlock(i) for i in range(num_blocks)]
self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_blocks)]
self.transformer_blocks = [JointTransformerBlock(i) for i in range(num_transformer_blocks)]
self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_transformer_blocks)]
zero_init = nn.init.constant(0)
self.controlnet_x_embedder = nn.Linear(64, 3072).apply(zero_init)
self.controlnet_blocks = [nn.Linear(3072, 3072).apply(zero_init) for _ in range(num_blocks)]
self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_blocks)]
self.controlnet_x_embedder = nn.Linear(64, 3072).apply(nn.init.constant(0))
self.controlnet_blocks = [nn.Linear(3072, 3072).apply(nn.init.constant(0)) for _ in range(num_transformer_blocks)] # fmt: off
self.controlnet_single_blocks = [nn.Linear(3072, 3072) for _ in range(num_single_transformer_blocks)]
def __call__(
self,
t: int,
config: RuntimeConfig,
hidden_states: mx.array,
prompt_embeds: mx.array,
pooled_prompt_embeds: mx.array,
hidden_states: mx.array,
controlnet_cond: mx.array,
config: RuntimeConfig,
controlnet_condition: mx.array,
) -> (list[mx.array], list[mx.array]):
time_step = config.sigmas[t] * config.num_train_steps
time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
hidden_states = self.x_embedder(hidden_states)
hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond)
conditioning_scale = config.config.controlnet_strength
guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision)
text_embeddings = self.time_text_embed(time_step, pooled_prompt_embeds, guidance)
# 1. Create embeddings
hidden_states = self.x_embedder(hidden_states) + self.controlnet_x_embedder(controlnet_condition)
encoder_hidden_states = self.context_embedder(prompt_embeds)
txt_ids = Transformer.prepare_text_ids(seq_len=prompt_embeds.shape[1])
img_ids = Transformer.prepare_latent_image_ids(height=config.height, width=config.width)
ids = mx.concatenate((txt_ids, img_ids), axis=1)
image_rotary_emb = self.pos_embed(ids)
text_embeddings = Transformer.compute_text_embeddings(t, pooled_prompt_embeds, self.time_text_embed, config)
image_rotary_embeddings = Transformer.compute_rotary_embeddings(prompt_embeds, self.pos_embed, config)
block_samples = ()
for block in self.transformer_blocks:
encoder_hidden_states, hidden_states = block(
# 2. Run the joint transformer blocks
controlnet_block_samples = []
for idx, block in enumerate(self.transformer_blocks):
encoder_hidden_states, hidden_states = self._apply_joint_transformer_block(
idx=idx,
block=block,
config=config,
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb,
image_rotary_embeddings=image_rotary_embeddings,
controlnet_block_samples=controlnet_block_samples,
)
block_samples = block_samples + (hidden_states,)
# 3. Concat the hidden states
hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
# controlnet block
controlnet_block_samples = ()
for block_sample, controlnet_block in zip(block_samples, self.controlnet_blocks):
block_sample = controlnet_block(block_sample)
controlnet_block_samples = controlnet_block_samples + (block_sample,)
single_block_samples = ()
for block in self.single_transformer_blocks:
hidden_states = block(
# 4. Run the single transformer blocks
controlnet_single_block_samples = []
for idx, block in enumerate(self.single_transformer_blocks):
hidden_states = self._apply_single_transformer_block(
idx=idx,
block=block,
config=config,
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb,
image_rotary_embeddings=image_rotary_embeddings,
controlnet_single_block_samples=controlnet_single_block_samples,
)
single_block_samples = single_block_samples + (hidden_states[:, encoder_hidden_states.shape[1] :],)
controlnet_single_block_samples = ()
for single_block_sample, controlnet_block in zip(single_block_samples, self.controlnet_single_blocks):
single_block_sample = controlnet_block(single_block_sample)
controlnet_single_block_samples = controlnet_single_block_samples + (single_block_sample,)
# scaling
controlnet_block_samples = [sample * conditioning_scale for sample in controlnet_block_samples]
controlnet_single_block_samples = [sample * conditioning_scale for sample in controlnet_single_block_samples]
return controlnet_block_samples, controlnet_single_block_samples
def _apply_single_transformer_block(
self,
idx: int,
config: RuntimeConfig,
block: SingleTransformerBlock,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
text_embeddings: mx.array,
image_rotary_embeddings: mx.array,
controlnet_single_block_samples: list[mx.array],
) -> mx.array:
# 1. Apply single transformer block
hidden_states = block(
hidden_states=hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_embeddings,
)
# 2. Apply controlnet block
states = hidden_states[:, encoder_hidden_states.shape[1] :]
controlnet_sample = self.controlnet_single_blocks[idx](states)
scaled_controlnet_sample = controlnet_sample * config.config.controlnet_strength
controlnet_single_block_samples.append(scaled_controlnet_sample)
return hidden_states
def _apply_joint_transformer_block(
self,
idx: int,
config: RuntimeConfig,
block: JointTransformerBlock,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
text_embeddings: mx.array,
image_rotary_embeddings: mx.array,
controlnet_block_samples: list[mx.array],
) -> tuple[mx.array, mx.array]:
# 1. Apply joint transformer block
encoder_hidden_states, hidden_states = block(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_embeddings,
)
# 2. Apply controlnet block
controlnet_sample = self.controlnet_blocks[idx](hidden_states)
scaled_controlnet_example = controlnet_sample * config.config.controlnet_strength
controlnet_block_samples.append(scaled_controlnet_example)
return encoder_hidden_states, hidden_states

12
src/mflux/controlnet/weight_handler_controlnet.py
View File

@ -8,6 +8,8 @@ from mlx.utils import tree_unflatten
from mflux.weights.weight_handler import MetaData
from mflux.weights.weight_util import WeightUtil
CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
class WeightHandlerControlnet:
def __init__(self, meta_data: MetaData, config: dict, controlnet_transformer: dict | None = None):
@ -16,8 +18,8 @@ class WeightHandlerControlnet:
self.config = config
@staticmethod
def load_controlnet_transformer(controlnet_id: str) -> "WeightHandlerControlnet":
controlnet_path = Path(snapshot_download(repo_id=controlnet_id, allow_patterns=["*.safetensors", "config.json"])) # fmt:off
def load_controlnet_transformer() -> "WeightHandlerControlnet":
controlnet_path = Path(snapshot_download(repo_id=CONTROLNET_ID, allow_patterns=["*.safetensors", "config.json"])) # fmt:off
file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
weights = list(mx.load(str(file)).items())
@ -60,3 +62,9 @@ class WeightHandlerControlnet:
controlnet_transformer=weights,
meta_data=MetaData(quantization_level=quantization_level)
) # fmt:off
def num_transformer_blocks(self) -> int:
return self.config["num_layers"]
def num_single_transformer_blocks(self) -> int:
return self.config["num_single_layers"]