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

Refactor: Transformer & Flux

Browse Source
This commit is contained in:
filipstrand 2025-02-01 20:57:08 +01:00
parent fa99cf0cc8
commit 0d9a5d626f
4 changed files with 127 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/mflux/dreambooth/optimization/dreambooth_loss.py
View File

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

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

@ -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

149
src/mflux/models/transformer/transformer.py
View File

@ -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]

3
src/mflux/weights/weight_handler.py
View File

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