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