98 lines
4.3 KiB
Python
98 lines
4.3 KiB
Python
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:
|
|
hidden_states = block.forward(
|
|
hidden_states=hidden_states,
|
|
text_embeddings=text_embeddings,
|
|
rotary_embeddings=image_rotary_emb,
|
|
)
|
|
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
|