diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index b961514..eda6a4b 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -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") diff --git a/src/mflux/controlnet/transformer_controlnet.py b/src/mflux/controlnet/transformer_controlnet.py index 6d9d4ad..8558965 100644 --- a/src/mflux/controlnet/transformer_controlnet.py +++ b/src/mflux/controlnet/transformer_controlnet.py @@ -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 diff --git a/src/mflux/controlnet/weight_handler_controlnet.py b/src/mflux/controlnet/weight_handler_controlnet.py index 0cdfd70..f8b80c8 100644 --- a/src/mflux/controlnet/weight_handler_controlnet.py +++ b/src/mflux/controlnet/weight_handler_controlnet.py @@ -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"] diff --git a/src/mflux/dreambooth/optimization/dreambooth_loss.py b/src/mflux/dreambooth/optimization/dreambooth_loss.py index c985833..dbc23e7 100644 --- a/src/mflux/dreambooth/optimization/dreambooth_loss.py +++ b/src/mflux/dreambooth/optimization/dreambooth_loss.py @@ -53,12 +53,12 @@ class DreamBoothLoss: ) # fmt: off # Predict the noise from timestep t - predicted_noise = flux.transformer.predict( + predicted_noise = flux.transformer( t=t, + config=config, + hidden_states=latents_t, prompt_embeds=example.prompt_embeds, pooled_prompt_embeds=example.pooled_prompt_embeds, - hidden_states=latents_t, - config=config, ) # Construct the loss (derivation in src/mflux/dreambooth/optimization/_loss_derivation) diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index fc6b8d6..388e224 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -51,7 +51,7 @@ class Flux1(nn.Module): # 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() @@ -100,12 +100,12 @@ class Flux1(nn.Module): for gen_step, t in enumerate(time_steps, 1): try: # 3.t Predict the noise - noise = self.transformer.predict( + 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, ) # 4.t Take one denoise step diff --git a/src/mflux/models/transformer/ada_layer_norm_zero.py b/src/mflux/models/transformer/ada_layer_norm_zero.py index 98f6173..bef933c 100644 --- a/src/mflux/models/transformer/ada_layer_norm_zero.py +++ b/src/mflux/models/transformer/ada_layer_norm_zero.py @@ -8,7 +8,7 @@ class AdaLayerNormZero(nn.Module): self.linear = nn.Linear(3072, 18432) self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False) - def __call__(self, x: mx.array, text_embeddings: mx.array): + def __call__(self, hidden_states: mx.array, text_embeddings: mx.array): text_embeddings = self.linear(nn.silu(text_embeddings)) chunk_size = 18432 // 6 shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size] @@ -17,5 +17,5 @@ class AdaLayerNormZero(nn.Module): shift_mlp = text_embeddings[:, 3 * chunk_size : 4 * chunk_size] scale_mlp = text_embeddings[:, 4 * chunk_size : 5 * chunk_size] gate_mlp = text_embeddings[:, 5 * chunk_size : 6 * chunk_size] - x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] - return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp diff --git a/src/mflux/models/transformer/ada_layer_norm_zero_single.py b/src/mflux/models/transformer/ada_layer_norm_zero_single.py index 425543f..3e254a6 100644 --- a/src/mflux/models/transformer/ada_layer_norm_zero_single.py +++ b/src/mflux/models/transformer/ada_layer_norm_zero_single.py @@ -8,11 +8,11 @@ class AdaLayerNormZeroSingle(nn.Module): self.linear = nn.Linear(3072, 3 * 3072) self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False) - def __call__(self, x: mx.array, text_embeddings: mx.array): + def __call__(self, hidden_states: mx.array, text_embeddings: mx.array) -> mx.array: text_embeddings = self.linear(nn.silu(text_embeddings)) chunk_size = 9216 // 3 shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size] scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size] gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size] - x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] - return x, gate_msa + hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return hidden_states, gate_msa diff --git a/src/mflux/models/transformer/joint_attention.py b/src/mflux/models/transformer/joint_attention.py index c114ae4..1c71005 100644 --- a/src/mflux/models/transformer/joint_attention.py +++ b/src/mflux/models/transformer/joint_attention.py @@ -30,6 +30,7 @@ class JointAttention(nn.Module): encoder_hidden_states: mx.array, image_rotary_emb: mx.array, ) -> (mx.array, mx.array): + # 1a. Compute Q,K,V for hidden_states query, key, value = AttentionUtils.process_qkv( hidden_states=hidden_states, to_q=self.to_q, @@ -40,6 +41,8 @@ class JointAttention(nn.Module): num_heads=self.num_heads, head_dim=self.head_dimension, ) + + # 1b. Compute Q,K,V for encoder_hidden_states enc_query, enc_key, enc_value = AttentionUtils.process_qkv( hidden_states=encoder_hidden_states, to_q=self.add_q_proj, @@ -51,12 +54,15 @@ class JointAttention(nn.Module): head_dim=self.head_dimension, ) + # 1c. Concatenate results query = mx.concatenate([enc_query, query], axis=2) key = mx.concatenate([enc_key, key], axis=2) value = mx.concatenate([enc_value, value], axis=2) + # 1d. Apply rope to Q,K query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=image_rotary_emb) + # 2. Compute attention hidden_states = AttentionUtils.compute_attention( query=query, key=key, @@ -66,12 +72,13 @@ class JointAttention(nn.Module): head_dim=self.head_dimension, ) + # 3. Separate the results encoder_hidden_states, hidden_states = ( hidden_states[:, : encoder_hidden_states.shape[1]], hidden_states[:, encoder_hidden_states.shape[1] :], ) + # 4. Project the output hidden_states = self.to_out[0](hidden_states) encoder_hidden_states = self.to_add_out(encoder_hidden_states) - return hidden_states, encoder_hidden_states diff --git a/src/mflux/models/transformer/joint_transformer_block.py b/src/mflux/models/transformer/joint_transformer_block.py index b89b4e6..2eb5419 100644 --- a/src/mflux/models/transformer/joint_transformer_block.py +++ b/src/mflux/models/transformer/joint_transformer_block.py @@ -11,12 +11,12 @@ class JointTransformerBlock(nn.Module): super().__init__() self.layer = layer self.norm1 = AdaLayerNormZero() - self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False) - self.ff = FeedForward(activation_function=nn.gelu) - self.attn = JointAttention() self.norm1_context = AdaLayerNormZero() - self.ff_context = FeedForward(activation_function=nn.gelu_approx) + self.attn = JointAttention() + self.norm2 = nn.LayerNorm(dims=3072, eps=1e-6, affine=False) self.norm2_context = nn.LayerNorm(dims=1536, eps=1e-6, affine=False) + self.ff = FeedForward(activation_function=nn.gelu) + self.ff_context = FeedForward(activation_function=nn.gelu_approx) def __call__( self, @@ -25,30 +25,67 @@ class JointTransformerBlock(nn.Module): text_embeddings: mx.array, rotary_embeddings: mx.array, ) -> (mx.array, mx.array): - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, text_embeddings) + # 1a. Compute norm for hidden_states + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states=hidden_states, + text_embeddings=text_embeddings + ) # fmt: off + # 1b. Compute norm for encoder_hidden_states norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( - x=encoder_hidden_states, text_embeddings=text_embeddings - ) + hidden_states=encoder_hidden_states, + text_embeddings=text_embeddings + ) # fmt: off + # 2. Compute attention attn_output, context_attn_output = self.attn( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, image_rotary_emb=rotary_embeddings, ) + # 3a. Apply norm and feed forward for hidden states + hidden_states = JointTransformerBlock._apply_norm_and_feed_forward( + hidden_states=hidden_states, + attn_output=attn_output, + gate_mlp=gate_mlp, + gate_msa=gate_msa, + scale_mlp=scale_mlp, + shift_mlp=shift_mlp, + norm_layer=self.norm2, + ff_layer=self.ff, + ) + + # 3b. Apply norm and feed forward for encoder hidden states + encoder_hidden_states = JointTransformerBlock._apply_norm_and_feed_forward( + hidden_states=encoder_hidden_states, + attn_output=context_attn_output, + gate_mlp=c_gate_mlp, + gate_msa=c_gate_msa, + scale_mlp=c_scale_mlp, + shift_mlp=c_shift_mlp, + norm_layer=self.norm2_context, + ff_layer=self.ff_context, + ) + + return encoder_hidden_states, hidden_states + + @staticmethod + def _apply_norm_and_feed_forward( + hidden_states: mx.array, + attn_output: mx.array, + gate_mlp: mx.array, + gate_msa: mx.array, + scale_mlp: mx.array, + shift_mlp: mx.array, + norm_layer: nn.Module, + ff_layer: nn.Module, + ) -> mx.array: attn_output = mx.expand_dims(gate_msa, axis=1) * attn_output hidden_states = hidden_states + attn_output - norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_layer(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - ff_output = self.ff(norm_hidden_states) + ff_output = ff_layer(norm_hidden_states) ff_output = mx.expand_dims(gate_mlp, axis=1) * ff_output hidden_states = hidden_states + ff_output - - context_attn_output = mx.expand_dims(c_gate_msa, axis=1) * context_attn_output - encoder_hidden_states = encoder_hidden_states + context_attn_output - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] - context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + mx.expand_dims(c_gate_mlp, axis=1) * context_ff_output - return encoder_hidden_states, hidden_states + return hidden_states diff --git a/src/mflux/models/transformer/single_block_attention.py b/src/mflux/models/transformer/single_block_attention.py index 88526b1..a666850 100644 --- a/src/mflux/models/transformer/single_block_attention.py +++ b/src/mflux/models/transformer/single_block_attention.py @@ -18,6 +18,7 @@ class SingleBlockAttention(nn.Module): self.norm_k = nn.RMSNorm(128) def __call__(self, hidden_states: mx.array, image_rotary_emb: mx.array) -> mx.array: + # 1a. Compute Q,K,V for hidden_states query, key, value = AttentionUtils.process_qkv( hidden_states=hidden_states, to_q=self.to_q, @@ -29,8 +30,10 @@ class SingleBlockAttention(nn.Module): head_dim=self.head_dimension, ) + # 1b. Apply rope to Q,K query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=image_rotary_emb) + # 2. Compute attention return AttentionUtils.compute_attention( query=query, key=key, diff --git a/src/mflux/models/transformer/single_transformer_block.py b/src/mflux/models/transformer/single_transformer_block.py index 1bab2a4..a2344ce 100644 --- a/src/mflux/models/transformer/single_transformer_block.py +++ b/src/mflux/models/transformer/single_transformer_block.py @@ -12,8 +12,8 @@ class SingleTransformerBlock(nn.Module): super().__init__() self.layer = layer self.norm = AdaLayerNormZeroSingle() - self.proj_mlp = nn.Linear(3072, 4 * 3072) self.attn = SingleBlockAttention() + self.proj_mlp = nn.Linear(3072, 4 * 3072) self.proj_out = nn.Linear(3072 + 4 * 3072, 3072) def __call__( @@ -22,15 +22,39 @@ class SingleTransformerBlock(nn.Module): text_embeddings: mx.array, rotary_embeddings: mx.array, ) -> (mx.array, mx.array): + # 0. Establish residual connection residual = hidden_states - norm_hidden_states, gate = self.norm(x=hidden_states, text_embeddings=text_embeddings) - mlp_hidden_states = nn.gelu_approx(self.proj_mlp(norm_hidden_states)) + + # 1. Compute norm for hidden_states + norm_hidden_states, gate = self.norm( + hidden_states=hidden_states, + text_embeddings=text_embeddings + ) # fmt: off + + # 2. Compute attention attn_output = self.attn( hidden_states=norm_hidden_states, image_rotary_emb=rotary_embeddings, ) + + # 3. Apply norm and feed forward for hidden states + hidden_states = self._apply_feed_forward_and_projection( + norm_hidden_states=norm_hidden_states, + attn_output=attn_output, + gate=gate, + ) + + return residual + hidden_states + + def _apply_feed_forward_and_projection( + self, + norm_hidden_states: mx.array, + attn_output: mx.array, + gate: mx.array, + ) -> mx.array: + feed_forward = self.proj_mlp(norm_hidden_states) + mlp_hidden_states = nn.gelu_approx(feed_forward) hidden_states = mx.concatenate([attn_output, mlp_hidden_states], axis=2) gate = mx.expand_dims(gate, axis=1) hidden_states = gate * self.proj_out(hidden_states) - hidden_states = residual + hidden_states return hidden_states diff --git a/src/mflux/models/transformer/transformer.py b/src/mflux/models/transformer/transformer.py index d0e5f23..0a20375 100644 --- a/src/mflux/models/transformer/transformer.py +++ b/src/mflux/models/transformer/transformer.py @@ -19,74 +19,141 @@ from mflux.models.transformer.time_text_embed import TimeTextEmbed class Transformer(nn.Module): - def __init__(self, model_config: ModelConfig, num_transformer_blocks: int): + def __init__( + self, + model_config: ModelConfig, + num_transformer_blocks: int = 19, + num_single_transformer_blocks: int = 38, + ): 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_transformer_blocks)] - self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)] + self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(num_single_transformer_blocks)] self.norm_out = AdaLayerNormContinuous(3072, 3072) self.proj_out = nn.Linear(3072, 64) - def predict( + def __call__( self, t: int, + config: RuntimeConfig, + hidden_states: mx.array, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, - hidden_states: mx.array, - config: RuntimeConfig, controlnet_block_samples: list[mx.array] | None = None, controlnet_single_block_samples: list[mx.array] | None = None, ) -> mx.array: - time_step = config.sigmas[t] * config.num_train_steps - time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) + # 1. Create embeddings hidden_states = self.x_embedder(hidden_states) - 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) 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(config.height, 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) + # 2. Run the joint transformer blocks for idx, block in enumerate(self.transformer_blocks): - encoder_hidden_states, hidden_states = block( + encoder_hidden_states, hidden_states = self._apply_joint_transformer_block( + idx=idx, + block=block, 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, ) - if controlnet_block_samples is not None and len(controlnet_block_samples) > 0: - interval_control = len(self.transformer_blocks) / len(controlnet_block_samples) - interval_control = int(math.ceil(interval_control)) - hidden_states = hidden_states + controlnet_block_samples[idx // interval_control] + # 3. Concat the hidden states hidden_states = mx.concatenate([encoder_hidden_states, hidden_states], axis=1) + # 4. Run the single transformer blocks for idx, block in enumerate(self.single_transformer_blocks): - hidden_states = block( + hidden_states = self._apply_single_transformer_block( + idx=idx, + block=block, 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, ) - if controlnet_single_block_samples is not None and len(controlnet_single_block_samples) > 0: - interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples) - interval_control = int(math.ceil(interval_control)) - hidden_states[:, encoder_hidden_states.shape[1] :, ...] = ( - hidden_states[:, encoder_hidden_states.shape[1] :, ...] - + controlnet_single_block_samples[idx // interval_control] - ) + # 5. Project the final output hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] hidden_states = self.norm_out(hidden_states, text_embeddings) hidden_states = self.proj_out(hidden_states) - noise = hidden_states - return noise + return hidden_states + + def _apply_single_transformer_block( + self, + idx: int, + 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 previously calculated controlnet result (if applicable) + sample = Transformer._get_controlnet_sample(idx, self.single_transformer_blocks, controlnet_single_block_samples) # fmt: off + hidden_states[:, encoder_hidden_states.shape[1] :, ...] += sample if sample is not None else 0 + + return hidden_states + + def _apply_joint_transformer_block( + self, + idx: int, + 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], + ) -> 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 previously calculated controlnet result (if applicable) + sample = Transformer._get_controlnet_sample(idx, self.transformer_blocks, controlnet_block_samples) + hidden_states += sample if sample is not None else 0 + + return encoder_hidden_states, hidden_states @staticmethod - def prepare_latent_image_ids(height: int, width: int) -> mx.array: + def compute_rotary_embeddings(prompt_embeds: mx.array, pos_embed: EmbedND, config: RuntimeConfig) -> mx.array: + 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 = pos_embed(ids) + return image_rotary_emb + + @staticmethod + def compute_text_embeddings( + t: int, + pooled_prompt_embeds: mx.array, + time_text_embed: TimeTextEmbed, + config: RuntimeConfig, + ) -> mx.array: + time_step = config.sigmas[t] * config.num_train_steps + time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) + guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision) + text_embeddings = time_text_embed(time_step, pooled_prompt_embeds, guidance) + return text_embeddings + + @staticmethod + def _prepare_latent_image_ids(height: int, width: int) -> mx.array: latent_width = width // 16 latent_height = height // 16 latent_image_ids = mx.zeros((latent_height, latent_width, 3)) @@ -97,5 +164,23 @@ class Transformer(nn.Module): return latent_image_ids @staticmethod - def prepare_text_ids(seq_len: mx.array) -> mx.array: + def _prepare_text_ids(seq_len: mx.array) -> mx.array: return mx.zeros((1, seq_len, 3)) + + @staticmethod + def _get_controlnet_sample( + idx: int, + blocks: mx.array, + controlnet_samples: list[mx.array] | None, + ) -> mx.array | None: # fmt: off + if controlnet_samples is None: + return None + + if len(controlnet_samples) == 0: + return None + + num_blocks = len(blocks) + num_samples = len(controlnet_samples) + interval_control = int(math.ceil(num_blocks / num_samples)) + control_index = idx // interval_control + return controlnet_samples[control_index] diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index 74db4ce..9679bb4 100644 --- a/src/mflux/weights/weight_handler.py +++ b/src/mflux/weights/weight_handler.py @@ -60,6 +60,9 @@ class WeightHandler: def num_transformer_blocks(self) -> int: return len(self.transformer["transformer_blocks"]) + def num_single_transformer_blocks(self) -> int: + return len(self.transformer["single_transformer_blocks"]) + @staticmethod def _load_clip_encoder(root_path: Path) -> (dict, int): weights, quantization_level, _ = WeightHandler._get_weights("text_encoder", root_path)