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

Fix formatting and some warnings

Browse Source
This commit is contained in:
filipstrand 2024-09-17 06:41:05 +02:00
parent 25e375d45d
commit 78bbf3d8fb
2 changed files with 13 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -230,6 +230,7 @@ class Flux1Controlnet:
ControlNetOutput = Tuple[list[mx.array], list[mx.array]]
class TransformerControlnet(nn.Module):
def __init__(
@ -270,8 +271,8 @@ class TransformerControlnet(nn.Module):
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)
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)
@ -293,7 +294,6 @@ class TransformerControlnet(nn.Module):
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(

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

@ -12,7 +12,6 @@ from mflux.models.transformer.single_transformer_block import SingleTransformerB
from mflux.models.transformer.time_text_embed import TimeTextEmbed
class Transformer(nn.Module):
def __init__(self, model_config: ModelConfig):
@ -33,8 +32,8 @@ class Transformer(nn.Module):
pooled_prompt_embeds: mx.array,
hidden_states: mx.array,
config: RuntimeConfig,
controlnet_block_samples: Tuple[mx.array] | None = None,
controlnet_single_block_samples: Tuple[mx.array] | None = None,
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)
@ -42,8 +41,8 @@ class Transformer(nn.Module):
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)
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)
@ -82,7 +81,7 @@ class Transformer(nn.Module):
return noise
@staticmethod
def _prepare_latent_image_ids(height: int, width: int) -> mx.array:
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))
@ -93,5 +92,5 @@ 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))