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

Move new controlnet module to its own file

Browse Source
This commit is contained in:
filipstrand 2024-09-17 06:50:02 +02:00
parent ad6f4ca8af
commit fc8099e908
2 changed files with 97 additions and 90 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,6 +1,5 @@
import logging
from pathlib import Path
from typing import Tuple
import PIL.Image
import mlx.core as mx
@ -12,12 +11,9 @@ from mflux.config.config import ConfigControlnet
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.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
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.time_text_embed import TimeTextEmbed
from mflux.models.transformer.transformer import Transformer
from mflux.models.vae.vae import VAE
from mflux.post_processing.image import GeneratedImage
@ -228,88 +224,3 @@ class Flux1Controlnet:
_save_weights(self.transformer_controlnet, "transformer_controlnet")
ControlNetOutput = Tuple[list[mx.array], list[mx.array]]
class TransformerControlnet(nn.Module):
def __init__(
self,
model_config: ModelConfig,
num_blocks: int,
num_single_blocks: int,
):
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)]
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)]
def forward(
self,
t: int,
prompt_embeds: mx.array,
pooled_prompt_embeds: mx.array,
hidden_states: mx.array,
controlnet_cond: mx.array,
config: RuntimeConfig,
) -> ControlNetOutput:
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.forward(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.forward(ids)
block_samples = ()
for block in self.transformer_blocks:
encoder_hidden_states, hidden_states = block.forward(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb
)
block_samples = block_samples + (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:
ctrlnet_hidden_states = block.forward(
hidden_states=ctrlnet_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb
)
single_block_samples = single_block_samples + (ctrlnet_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

96
src/mflux/controlnet/transformer_controlnet.py Normal file
View File

@ -0,0 +1,96 @@
from typing import Tuple
import mlx.core as mx
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.time_text_embed import TimeTextEmbed
from mflux.models.transformer.transformer import Transformer
class TransformerControlnet(nn.Module):
def __init__(
self,
model_config: ModelConfig,
num_blocks: int,
num_single_blocks: int,
):
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)]
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)]
def forward(
self,
t: int,
prompt_embeds: mx.array,
pooled_prompt_embeds: mx.array,
hidden_states: mx.array,
controlnet_cond: mx.array,
config: RuntimeConfig,
) -> (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.forward(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.forward(ids)
block_samples = ()
for block in self.transformer_blocks:
encoder_hidden_states, hidden_states = block.forward(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb
)
block_samples = block_samples + (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:
ctrlnet_hidden_states = block.forward(
hidden_states=ctrlnet_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb
)
single_block_samples = single_block_samples + (ctrlnet_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