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

Merge pull request #58 from anthonywu/intro-ruff-config

Browse Source 
introduce standardized `ruff check` and `ruff format` config
This commit is contained in:
Filip Strand 2024-09-21 23:25:57 +02:00 committed by GitHub
commit 86da9b2e43
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
68 changed files with 430 additions and 354 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
.pre-commit-config.yaml Normal file
View File

@ -0,0 +1,12 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.6.6
hooks:
# Run the linter.
- id: ruff
types_or: [python, pyi]
args: [--fix]
# Run the formatter.
- id: ruff-format
types_or: [python, pyi]

61
pyproject.toml
View File

@ -46,3 +46,64 @@ mflux-generate-controlnet = "mflux.generate_controlnet:main"
[tool.setuptools.packages.find]
where = ["src"]
include = ["mflux*"]
[tool.ruff]
line-length = 120
indent-width = 4
target-version = "py310"
respect-gitignore = true
exclude = [
"src/mflux/generate.py",
"src/mflux/generate_controlnet.py",
"src/mflux/save.py",
"src/mflux/flux/flux.py",
"src/mflux/controlnet/flux_controlnet.py",
"src/mflux/models/transformer/joint_transformer_block.py",
"src/mflux/models/vae/common/resnet_block_2d.py",
"src/mflux/models/vae/common/unet_mid_block.py",
"src/mflux/models/vae/encoder/down_block_2.py",
"src/mflux/models/vae/encoder/down_block_3.py",
"src/mflux/models/vae/decoder/up_block_3.py",
"src/mflux/models/vae/decoder/up_block_4.py"
]
[tool.ruff.lint]
# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default.
# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
# McCabe complexity (`C901`) by default.
select = ["E4", "E7", "E9", "F"]
ignore = []
# Allow fix for all enabled rules (when `--fix`) is provided.
fixable = ["ALL"]
unfixable = []
# Allow unused variables when underscore-prefixed.
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
[tool.ruff.format]
# Like Black, use double quotes for strings.
quote-style = "double"
# Like Black, indent with spaces, rather than tabs.
indent-style = "space"
# Like Black, respect magic trailing commas.
skip-magic-trailing-comma = false
# Like Black, automatically detect the appropriate line ending.
line-ending = "auto"
# Enable auto-formatting of code examples in docstrings. Markdown,
# reStructuredText code/literal blocks and doctests are all supported.
#
# This is currently disabled by default, but it is planned for this
# to be opt-out in the future.
docstring-code-format = false
# Set the line length limit used when formatting code snippets in
# docstrings.
#
# This only has an effect when the `docstring-code-format` setting is
# enabled.
docstring-code-line-length = "dynamic"

9
src/mflux/__init__.py
View File

@ -5,4 +5,11 @@ from mflux.controlnet.flux_controlnet import Flux1Controlnet
from mflux.flux.flux import Flux1
from mflux.post_processing.image_util import ImageUtil
__all__ = ["Flux1", "Flux1Controlnet", "Config", "ConfigControlnet", "ModelConfig", "ImageUtil"]
__all__ = [
"Flux1",
"Flux1Controlnet",
"Config",
"ConfigControlnet",
"ModelConfig",
"ImageUtil",
]

22
src/mflux/config/config.py
View File

@ -9,11 +9,11 @@ class Config:
precision: mx.Dtype = mx.bfloat16
def __init__(
self,
num_inference_steps: int = 4,
width: int = 1024,
height: int = 1024,
guidance: float = 4.0,
self,
num_inference_steps: int = 4,
width: int = 1024,
height: int = 1024,
guidance: float = 4.0,
):
if width % 16 != 0 or height % 16 != 0:
log.warning("Width and height should be multiples of 16. Rounding down.")
@ -25,12 +25,12 @@ class Config:
class ConfigControlnet(Config):
def __init__(
self,
num_inference_steps: int = 4,
width: int = 1024,
height: int = 1024,
guidance: float = 4.0,
controlnet_strength: float = 1.0,
self,
num_inference_steps: int = 4,
width: int = 1024,
height: int = 1024,
guidance: float = 4.0,
controlnet_strength: float = 1.0,
):
super().__init__(num_inference_steps, width, height, guidance)
self.controlnet_strength = controlnet_strength

10
src/mflux/config/model_config.py
View File

@ -6,11 +6,11 @@ class ModelConfig(Enum):
FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 1000, 256)
def __init__(
self,
model_name: str,
alias: str,
num_train_steps: int,
max_sequence_length: int,
self,
model_name: str,
alias: str,
num_train_steps: int,
max_sequence_length: int,
):
self.alias = alias
self.model_name = model_name

3
src/mflux/config/runtime_config.py
View File

@ -6,7 +6,6 @@ from mflux.config.model_config import ModelConfig
class RuntimeConfig:
def __init__(self, config: Config | ConfigControlnet, model_config: ModelConfig):
self.config = config
self.model_config = model_config
@ -35,7 +34,7 @@ class RuntimeConfig:
@property
def num_train_steps(self) -> int:
return self.model_config.num_train_steps
@property
def controlnet_strength(self) -> float:
if isinstance(self.config, ConfigControlnet):

1
src/mflux/controlnet/controlnet_util.py
View File

@ -4,7 +4,6 @@ import PIL
class ControlnetUtil:
@staticmethod
def preprocess_canny(img: PIL.Image) -> PIL.Image:
image_to_canny = np.array(img)

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

@ -29,13 +29,13 @@ CONTROLNET_ID = "InstantX/FLUX.1-dev-Controlnet-Canny"
class Flux1Controlnet:
def __init__(
self,
model_config: ModelConfig,
quantize: int | None = None,
local_path: str | None = None,
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
controlnet_path: str | None = None,
self,
model_config: ModelConfig,
quantize: int | None = None,
local_path: str | None = None,
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
controlnet_path: str | None = None,
):
self.lora_paths = lora_paths
self.lora_scales = lora_scales
@ -112,7 +112,7 @@ class Flux1Controlnet:
control_image = ControlnetUtil.preprocess_canny(control_image)
controlnet_cond = ImageUtil.to_array(control_image)
controlnet_cong = self.vae.encode(controlnet_cond)
# the rescaling in the next line is not in the huggingface code, but without it the images from
# the rescaling in the next line is not in the huggingface code, but without it the images from
# the chosen controlnet model are very bad
controlnet_cond = (controlnet_cong / self.vae.scaling_factor) + self.vae.shift_factor
controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
@ -158,7 +158,7 @@ class Flux1Controlnet:
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict['elapsed'],
generation_time=time_steps.format_dict["elapsed"],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
config=config,
@ -170,7 +170,7 @@ class Flux1Controlnet:
latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2))
return latents
@staticmethod
def _pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2))

37
src/mflux/controlnet/transformer_controlnet.py
View File

@ -4,19 +4,22 @@ 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.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,
self,
model_config: ModelConfig,
num_blocks: int,
num_single_blocks: int,
):
super().__init__()
self.pos_embed = EmbedND()
@ -33,13 +36,13 @@ class TransformerControlnet(nn.Module):
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,
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)
@ -61,7 +64,7 @@ class TransformerControlnet(nn.Module):
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb
rotary_embeddings=image_rotary_emb,
)
block_samples = block_samples + (hidden_states,)
@ -78,9 +81,9 @@ class TransformerControlnet(nn.Module):
hidden_states = block.forward(
hidden_states=hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb
rotary_embeddings=image_rotary_emb,
)
single_block_samples = single_block_samples + (hidden_states[:, encoder_hidden_states.shape[1]:],)
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):

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

@ -19,14 +19,13 @@ from mflux.weights.weight_handler import WeightHandler
class Flux1:
def __init__(
self,
model_config: ModelConfig,
quantize: int | None = None,
local_path: str | None = None,
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
self,
model_config: ModelConfig,
quantize: int | None = None,
local_path: str | None = None,
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
):
self.lora_paths = lora_paths
self.lora_scales = lora_scales
@ -110,7 +109,7 @@ class Flux1:
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict['elapsed'],
generation_time=time_steps.format_dict["elapsed"],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
config=config,

34
src/mflux/generate.py
View File

@ -5,20 +5,20 @@ from mflux import Flux1, Config, ModelConfig
def main():
parser = argparse.ArgumentParser(description='Generate an image based on a prompt.')
parser.add_argument('--prompt', type=str, required=True, help='The textual description of the image to generate.')
parser.add_argument('--output', type=str, default="image.png", help='The filename for the output image. Default is "image.png".')
parser.add_argument('--model', "-m", type=str, required=True, choices=["dev", "schnell"], help='The model to use ("schnell" or "dev").')
parser.add_argument('--seed', type=int, default=None, help='Entropy Seed (Default is time-based random-seed)')
parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)')
parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)')
parser.add_argument('--steps', type=int, default=None, help='Inference Steps')
parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)')
parser.add_argument('--quantize', "-q", type=int, choices=[4, 8], default=None, help='Quantize the model (4 or 8, Default is None)')
parser.add_argument('--path', type=str, default=None, help='Local path for loading a model from disk')
parser.add_argument('--lora-paths', type=str, nargs='*', default=None, help='Local safetensors for applying LORA from disk')
parser.add_argument('--lora-scales', type=float, nargs='*', default=None, help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.')
parser.add_argument('--metadata', action='store_true', help='Export image metadata as a JSON file.')
parser = argparse.ArgumentParser(description="Generate an image based on a prompt.")
parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.")
parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".")
parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").")
parser.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)")
parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)")
parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)")
parser.add_argument("--steps", type=int, default=None, help="Inference Steps")
parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)")
parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)")
parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk")
parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk")
parser.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.")
parser.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.")
args = parser.parse_args()
@ -34,7 +34,7 @@ def main():
quantize=args.quantize,
local_path=args.path,
lora_paths=args.lora_paths,
lora_scales=args.lora_scales
lora_scales=args.lora_scales,
)
# Generate an image
@ -46,12 +46,12 @@ def main():
height=args.height,
width=args.width,
guidance=args.guidance,
)
),
)
# Save the image
image.save(path=args.output, export_json_metadata=args.metadata)
if __name__ == '__main__':
if __name__ == "__main__":
main()

40
src/mflux/generate_controlnet.py
View File

@ -5,22 +5,22 @@ from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, ImageUtil
def main():
parser = argparse.ArgumentParser(description='Generate an image based on a prompt.')
parser.add_argument('--prompt', type=str, required=True, help='The textual description of the image to generate.')
parser.add_argument('--control-image-path', type=str, required=True, help='Local path of the image to use as input for controlnet.')
parser.add_argument('--output', type=str, default="image.png", help='The filename for the output image. Default is "image.png".')
parser.add_argument('--model', "-m", type=str, required=True, choices=["dev", "schnell"], help='The model to use ("schnell" or "dev").')
parser.add_argument('--seed', type=int, default=None, help='Entropy Seed (Default is time-based random-seed)')
parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)')
parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)')
parser.add_argument('--steps', type=int, default=None, help='Inference Steps')
parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)')
parser.add_argument('--controlnet-strength', type=float, default=0.7, help='Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.7)')
parser.add_argument('--quantize', "-q", type=int, choices=[4, 8], default=None, help='Quantize the model (4 or 8, Default is None)')
parser.add_argument('--path', type=str, default=None, help='Local path for loading a model from disk')
parser.add_argument('--lora-paths', type=str, nargs='*', default=None, help='Local safetensors for applying LORA from disk')
parser.add_argument('--lora-scales', type=float, nargs='*', default=None, help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.')
parser.add_argument('--metadata', action='store_true', help='Export image metadata as a JSON file.')
parser = argparse.ArgumentParser(description="Generate an image based on a prompt.")
parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.")
parser.add_argument("--control-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.")
parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".")
parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").")
parser.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)")
parser.add_argument("--height", type=int, default=1024, help="Image height (Default is 1024)")
parser.add_argument("--width", type=int, default=1024, help="Image width (Default is 1024)")
parser.add_argument("--steps", type=int, default=None, help="Inference Steps")
parser.add_argument("--guidance", type=float, default=3.5, help="Guidance Scale (Default is 3.5)")
parser.add_argument("--controlnet-strength", type=float, default=0.7, help="Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.7)")
parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=None, help="Quantize the model (4 or 8, Default is None)")
parser.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk")
parser.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk")
parser.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.")
parser.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.")
args = parser.parse_args()
@ -36,7 +36,7 @@ def main():
quantize=args.quantize,
local_path=args.path,
lora_paths=args.lora_paths,
lora_scales=args.lora_scales
lora_scales=args.lora_scales,
)
# Generate an image
@ -49,13 +49,13 @@ def main():
height=args.height,
width=args.width,
guidance=args.guidance,
controlnet_strength=args.controlnet_strength
)
controlnet_strength=args.controlnet_strength,
),
)
# Save the image
image.save(path=args.output, export_json_metadata=args.metadata)
if __name__ == '__main__':
if __name__ == "__main__":
main()

1
src/mflux/models/text_encoder/clip_encoder/clip_embeddings.py
View File

@ -5,7 +5,6 @@ from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
class CLIPEmbeddings(nn.Module):
def __init__(self, dims: int):
super().__init__()
self.position_embedding = nn.Embedding(num_embeddings=TokenizerCLIP.MAX_TOKEN_LENGTH, dims=dims)

1
src/mflux/models/text_encoder/clip_encoder/clip_encoder.py
View File

@ -5,7 +5,6 @@ from mflux.models.text_encoder.clip_encoder.clip_text_model import CLIPTextModel
class CLIPEncoder(nn.Module):
def __init__(self):
super().__init__()
self.text_model = CLIPTextModel(dims=768, num_encoder_layers=12)

1
src/mflux/models/text_encoder/clip_encoder/clip_encoder_layer.py
View File

@ -6,7 +6,6 @@ from mflux.models.text_encoder.clip_encoder.clip_sdpa_attention import CLIPSdpaA
class CLIPEncoderLayer(nn.Module):
def __init__(self, layer: int):
super().__init__()
self.self_attn = CLIPSdpaAttention()

1
src/mflux/models/text_encoder/clip_encoder/clip_mlp.py
View File

@ -3,7 +3,6 @@ from mlx import nn
class CLIPMLP(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(input_dims=768, output_dims=3072)

2
src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py
View File

@ -36,7 +36,7 @@ class CLIPSdpaAttention(nn.Module):
scores = (query * scale) @ key.transpose(0, 1, 3, 2)
scores = scores + mask
attn = mx.softmax(scores, axis=-1)
hidden_states = (attn @ value)
hidden_states = attn @ value
return hidden_states
@staticmethod

1
src/mflux/models/text_encoder/clip_encoder/clip_text_model.py
View File

@ -6,7 +6,6 @@ from mflux.models.text_encoder.clip_encoder.encoder_clip import EncoderCLIP
class CLIPTextModel(nn.Module):
def __init__(self, dims: int, num_encoder_layers: int):
super().__init__()
self.encoder = EncoderCLIP(num_encoder_layers)

10
src/mflux/models/text_encoder/clip_encoder/encoder_clip.py
View File

@ -1,11 +1,12 @@
import mlx.core as mx
from mlx import nn
from mflux.models.text_encoder.clip_encoder.clip_encoder_layer import CLIPEncoderLayer
from mflux.models.text_encoder.clip_encoder.clip_encoder_layer import (
CLIPEncoderLayer,
)
class EncoderCLIP(nn.Module):
def __init__(self, num_encoder_layers: int):
super().__init__()
self.layers = [CLIPEncoderLayer(i) for i in range(num_encoder_layers)]
@ -13,8 +14,5 @@ class EncoderCLIP(nn.Module):
def forward(self, tokens: mx.array, causal_attention_mask: mx.array) -> mx.array:
hidden_states = tokens
for encoder_layer in self.layers:
hidden_states = encoder_layer.forward(
hidden_states,
causal_attention_mask
)
hidden_states = encoder_layer.forward(hidden_states, causal_attention_mask)
return hidden_states

5
src/mflux/models/text_encoder/t5_encoder/t5_attention.py
View File

@ -2,11 +2,12 @@ import mlx.core as mx
from mlx import nn
from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
from mflux.models.text_encoder.t5_encoder.t5_self_attention import T5SelfAttention
from mflux.models.text_encoder.t5_encoder.t5_self_attention import (
T5SelfAttention,
)
class T5Attention(nn.Module):
def __init__(self):
super().__init__()
self.SelfAttention = T5SelfAttention()

1
src/mflux/models/text_encoder/t5_encoder/t5_block.py
View File

@ -6,7 +6,6 @@ from mflux.models.text_encoder.t5_encoder.t5_feed_forward import T5FeedForward
class T5Block(nn.Module):
def __init__(self, layer: int):
super().__init__()
self.attention = T5Attention()

7
src/mflux/models/text_encoder/t5_encoder/t5_dense_relu_dense.py
View File

@ -5,7 +5,6 @@ import mlx.core as mx
class T5DenseReluDense(nn.Module):
def __init__(self):
super().__init__()
self.wi_0 = nn.Linear(4096, 10240, bias=False)
@ -21,4 +20,8 @@ class T5DenseReluDense(nn.Module):
@staticmethod
def new_gelu(input_array: mx.array) -> mx.array:
return 0.5 * input_array * (1.0 + mx.tanh(math.sqrt(2.0 / math.pi) * (input_array + 0.044715 * mx.power(input_array, 3.0))))
return (
0.5
* input_array
* (1.0 + mx.tanh(math.sqrt(2.0 / math.pi) * (input_array + 0.044715 * mx.power(input_array, 3.0))))
)

1
src/mflux/models/text_encoder/t5_encoder/t5_encoder.py
View File

@ -6,7 +6,6 @@ from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
class T5Encoder(nn.Module):
def __init__(self):
super().__init__()
self.shared = nn.Embedding(num_embeddings=32128, dims=4096)

9
src/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py
View File

@ -1,14 +1,13 @@
import math
from mlx import nn
import mlx.core as mx
from mlx import nn
from mflux.models.text_encoder.t5_encoder.t5_dense_relu_dense import T5DenseReluDense
from mflux.models.text_encoder.t5_encoder.t5_dense_relu_dense import (
T5DenseReluDense,
)
from mflux.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
class T5FeedForward(nn.Module):
def __init__(self):
super().__init__()
self.layer_norm = T5LayerNorm()

7
src/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py
View File

@ -3,13 +3,16 @@ from mlx import nn
class T5LayerNorm(nn.Module):
def __init__(self):
super().__init__()
self.weight = mx.ones((4096,))
self.variance_epsilon = 1e-06
def forward(self, hidden_states: mx.array) -> mx.array:
variance = mx.mean(mx.power(hidden_states.astype(mx.float32), 2), axis=-1, keepdims=True)
variance = mx.mean(
mx.power(hidden_states.astype(mx.float32), 2),
axis=-1,
keepdims=True,
)
hidden_states = hidden_states * mx.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states

3
src/mflux/models/text_encoder/t5_encoder/t5_self_attention.py
View File

@ -5,7 +5,6 @@ from mlx import nn
class T5SelfAttention(nn.Module):
def __init__(self):
super().__init__()
self.q = nn.Linear(4096, 4096, bias=False)
@ -64,7 +63,7 @@ class T5SelfAttention(nn.Module):
).astype(mx.int32)
relative_position_if_large = mx.minimum(
relative_position_if_large,
mx.full(relative_position_if_large.shape, num_buckets - 1)
mx.full(relative_position_if_large.shape, num_buckets - 1),
)
relative_buckets += mx.where(is_small, relative_position, relative_position_if_large)

6
src/mflux/models/transformer/ada_layer_norm_continous.py
View File

@ -5,7 +5,6 @@ from mflux.config.config import Config
class AdaLayerNormContinuous(nn.Module):
def __init__(self, embedding_dim: int, conditioning_embedding_dim: int):
super().__init__()
self.embedding_dim = embedding_dim
@ -15,8 +14,7 @@ class AdaLayerNormContinuous(nn.Module):
def forward(self, x: mx.array, text_embeddings: mx.array) -> mx.array:
text_embeddings = self.linear(nn.silu(text_embeddings).astype(Config.precision))
chunk_size = self.embedding_dim
scale = text_embeddings[:, 0*chunk_size:1*chunk_size]
shift = text_embeddings[:, 1*chunk_size:2*chunk_size]
scale = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
shift = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
return x

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

@ -3,7 +3,6 @@ import mlx.core as mx
class AdaLayerNormZero(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(3072, 18432)
@ -12,11 +11,11 @@ class AdaLayerNormZero(nn.Module):
def forward(self, x: mx.array, text_embeddings: mx.array):
text_embeddings = self.linear(nn.silu(text_embeddings))
chunk_size = 18432 // 6
shift_msa = text_embeddings[:, 0*chunk_size:1*chunk_size]
scale_msa = text_embeddings[:, 1*chunk_size:2*chunk_size]
gate_msa = text_embeddings[:, 2*chunk_size:3*chunk_size]
shift_mlp = text_embeddings[:, 3*chunk_size:4*chunk_size]
scale_mlp = text_embeddings[:, 4*chunk_size:5*chunk_size]
gate_mlp = text_embeddings[:, 5*chunk_size:6*chunk_size]
shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size]
shift_mlp = text_embeddings[:, 3 * chunk_size : 4 * chunk_size]
scale_mlp = text_embeddings[:, 4 * chunk_size : 5 * chunk_size]
gate_mlp = text_embeddings[:, 5 * chunk_size : 6 * chunk_size]
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return x, gate_msa, shift_mlp, scale_mlp, gate_mlp

9
src/mflux/models/transformer/ada_layer_norm_zero_single.py
View File

@ -3,17 +3,16 @@ import mlx.core as mx
class AdaLayerNormZeroSingle(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(3072, 3*3072)
self.linear = nn.Linear(3072, 3 * 3072)
self.norm = nn.LayerNorm(dims=3072, eps=1e-6, affine=False)
def forward(self, x: mx.array, text_embeddings: mx.array):
text_embeddings = self.linear(nn.silu(text_embeddings))
chunk_size = 9216 // 3
shift_msa = text_embeddings[:, 0*chunk_size:1*chunk_size]
scale_msa = text_embeddings[:, 1*chunk_size:2*chunk_size]
gate_msa = text_embeddings[:, 2*chunk_size:3*chunk_size]
shift_msa = text_embeddings[:, 0 * chunk_size : 1 * chunk_size]
scale_msa = text_embeddings[:, 1 * chunk_size : 2 * chunk_size]
gate_msa = text_embeddings[:, 2 * chunk_size : 3 * chunk_size]
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return x, gate_msa

3
src/mflux/models/transformer/embed_nd.py
View File

@ -3,7 +3,6 @@ from mlx import nn
class EmbedND(nn.Module):
def __init__(self):
super().__init__()
self.dim = 3072
@ -20,7 +19,7 @@ class EmbedND(nn.Module):
@staticmethod
def rope(pos: mx.array, dim: int, theta: float) -> mx.array:
scale = mx.arange(0, dim, 2, dtype=mx.float32) / dim
omega = 1.0 / (theta ** scale)
omega = 1.0 / (theta**scale)
batch_size, seq_length = pos.shape
pos_expanded = mx.expand_dims(pos, axis=-1)
omega_expanded = mx.expand_dims(omega, axis=0)

1
src/mflux/models/transformer/feed_forward.py
View File

@ -3,7 +3,6 @@ import mlx.core as mx
class FeedForward(nn.Module):
def __init__(self, activation_function):
super().__init__()
self.linear1 = nn.Linear(3072, 12288)

1
src/mflux/models/transformer/guidance_embedder.py
View File

@ -3,7 +3,6 @@ import mlx.core as mx
class GuidanceEmbedder(nn.Module):
def __init__(self):
super().__init__()
self.linear_1 = nn.Linear(256, 3072)

34
src/mflux/models/transformer/joint_attention.py
View File

@ -23,13 +23,11 @@ class JointAttention(nn.Module):
self.norm_added_k = nn.RMSNorm(128)
def forward(
self,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
image_rotary_emb: mx.array
self,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
image_rotary_emb: mx.array,
) -> (mx.array, mx.array):
residual = hidden_states
query = self.to_q(hidden_states)
key = self.to_k(hidden_states)
value = self.to_v(hidden_states)
@ -45,9 +43,18 @@ class JointAttention(nn.Module):
encoder_hidden_states_key_proj = self.add_k_proj(encoder_hidden_states)
encoder_hidden_states_value_proj = self.add_v_proj(encoder_hidden_states)
encoder_hidden_states_query_proj = mx.transpose(mx.reshape(encoder_hidden_states_query_proj, (1, -1, 24, 128)), (0, 2, 1, 3))
encoder_hidden_states_key_proj = mx.transpose(mx.reshape(encoder_hidden_states_key_proj, (1, -1, 24, 128)), (0, 2, 1, 3))
encoder_hidden_states_value_proj = mx.transpose(mx.reshape(encoder_hidden_states_value_proj, (1, -1, 24, 128)), (0, 2, 1, 3))
encoder_hidden_states_query_proj = mx.transpose(
mx.reshape(encoder_hidden_states_query_proj, (1, -1, 24, 128)),
(0, 2, 1, 3),
)
encoder_hidden_states_key_proj = mx.transpose(
mx.reshape(encoder_hidden_states_key_proj, (1, -1, 24, 128)),
(0, 2, 1, 3),
)
encoder_hidden_states_value_proj = mx.transpose(
mx.reshape(encoder_hidden_states_value_proj, (1, -1, 24, 128)),
(0, 2, 1, 3),
)
encoder_hidden_states_query_proj = self.norm_added_q(encoder_hidden_states_query_proj)
encoder_hidden_states_key_proj = self.norm_added_k(encoder_hidden_states_key_proj)
@ -60,10 +67,13 @@ class JointAttention(nn.Module):
hidden_states = JointAttention.attention(query, key, value)
hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3))
hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension))
hidden_states = mx.reshape(
hidden_states,
(self.batch_size, -1, self.num_heads * self.head_dimension),
)
encoder_hidden_states, hidden_states = (
hidden_states[:, : encoder_hidden_states.shape[1]],
hidden_states[:, encoder_hidden_states.shape[1]:],
hidden_states[:, encoder_hidden_states.shape[1] :],
)
hidden_states = self.to_out[0](hidden_states)
@ -76,7 +86,7 @@ class JointAttention(nn.Module):
scale = 1 / mx.sqrt(query.shape[-1])
scores = (query * scale) @ key.transpose(0, 1, 3, 2)
attn = mx.softmax(scores, axis=-1)
hidden_states = (attn @ value)
hidden_states = attn @ value
return hidden_states
@staticmethod

11
src/mflux/models/transformer/joint_transformer_block.py
View File

@ -7,7 +7,6 @@ from mflux.models.transformer.joint_attention import JointAttention
class JointTransformerBlock(nn.Module):
def __init__(self, layer):
super().__init__()
self.layer = layer
@ -20,11 +19,11 @@ class JointTransformerBlock(nn.Module):
self.norm2_context = nn.LayerNorm(dims=1536, eps=1e-6, affine=False)
def forward(
self,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
text_embeddings: mx.array,
rotary_embeddings: mx.array
self,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
text_embeddings: mx.array,
rotary_embeddings: mx.array,
) -> (mx.array, mx.array):
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1.forward(hidden_states, text_embeddings)

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

@ -15,11 +15,7 @@ class SingleBlockAttention(nn.Module):
self.norm_q = nn.RMSNorm(128)
self.norm_k = nn.RMSNorm(128)
def forward(
self,
hidden_states: mx.array,
image_rotary_emb: mx.array
) -> (mx.array, mx.array):
def forward(self, hidden_states: mx.array, image_rotary_emb: mx.array) -> (mx.array, mx.array):
query = self.to_q(hidden_states)
key = self.to_k(hidden_states)
value = self.to_v(hidden_states)
@ -35,7 +31,10 @@ class SingleBlockAttention(nn.Module):
hidden_states = SingleBlockAttention.attention(query, key, value)
hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3))
hidden_states = mx.reshape(hidden_states, (self.batch_size, -1, self.num_heads * self.head_dimension))
hidden_states = mx.reshape(
hidden_states,
(self.batch_size, -1, self.num_heads * self.head_dimension),
)
return hidden_states
@ -44,7 +43,7 @@ class SingleBlockAttention(nn.Module):
scale = 1 / mx.sqrt(query.shape[-1])
scores = (query * scale) @ key.transpose(0, 1, 3, 2)
attn = mx.softmax(scores, axis=-1)
hidden_states = (attn @ value)
hidden_states = attn @ value
return hidden_states
@staticmethod

17
src/mflux/models/transformer/single_transformer_block.py
View File

@ -1,25 +1,26 @@
import mlx.core as mx
from mlx import nn
from mflux.models.transformer.ada_layer_norm_zero_single import AdaLayerNormZeroSingle
from mflux.models.transformer.ada_layer_norm_zero_single import (
AdaLayerNormZeroSingle,
)
from mflux.models.transformer.single_block_attention import SingleBlockAttention
class SingleTransformerBlock(nn.Module):
def __init__(self, layer):
super().__init__()
self.layer = layer
self.norm = AdaLayerNormZeroSingle()
self.proj_mlp = nn.Linear(3072, 4*3072)
self.proj_mlp = nn.Linear(3072, 4 * 3072)
self.attn = SingleBlockAttention()
self.proj_out = nn.Linear(3072 + 4*3072, 3072)
self.proj_out = nn.Linear(3072 + 4 * 3072, 3072)
def forward(
self,
hidden_states: mx.array,
text_embeddings: mx.array,
rotary_embeddings: mx.array
self,
hidden_states: mx.array,
text_embeddings: mx.array,
rotary_embeddings: mx.array,
) -> (mx.array, mx.array):
residual = hidden_states
norm_hidden_states, gate = self.norm.forward(x=hidden_states, text_embeddings=text_embeddings)

1
src/mflux/models/transformer/text_embedder.py
View File

@ -3,7 +3,6 @@ import mlx.core as mx
class TextEmbedder(nn.Module):
def __init__(self):
super().__init__()
self.linear_1 = nn.Linear(768, 3072)

9
src/mflux/models/transformer/time_text_embed.py
View File

@ -10,14 +10,18 @@ from mflux.models.transformer.guidance_embedder import GuidanceEmbedder
class TimeTextEmbed(nn.Module):
def __init__(self, model_config: ModelConfig):
super().__init__()
self.text_embedder = TextEmbedder()
self.guidance_embedder = GuidanceEmbedder() if model_config == ModelConfig.FLUX1_DEV else None
self.timestep_embedder = TimestepEmbedder()
def forward(self, time_step: mx.array, pooled_projection: mx.array, guidance: mx.array) -> mx.array:
def forward(
self,
time_step: mx.array,
pooled_projection: mx.array,
guidance: mx.array,
) -> mx.array:
time_steps_proj = self._time_proj(time_step)
time_steps_emb = self.timestep_embedder.forward(time_steps_proj)
if self.guidance_embedder is not None:
@ -37,4 +41,3 @@ class TimeTextEmbed(nn.Module):
emb = mx.concatenate([mx.sin(emb), mx.cos(emb)], axis=-1)
emb = mx.concatenate([emb[:, half_dim:], emb[:, :half_dim]], axis=-1)
return emb

1
src/mflux/models/transformer/timestep_embedder.py
View File

@ -3,7 +3,6 @@ import mlx.core as mx
class TimestepEmbedder(nn.Module):
def __init__(self):
super().__init__()
self.linear_1 = nn.Linear(256, 3072)

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

@ -5,15 +5,20 @@ from mlx import nn
from mflux.config.model_config import ModelConfig
from mflux.config.runtime_config import RuntimeConfig
from mflux.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
from mflux.models.transformer.ada_layer_norm_continous import (
AdaLayerNormContinuous,
)
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.joint_transformer_block import (
JointTransformerBlock,
)
from mflux.models.transformer.single_transformer_block import (
SingleTransformerBlock,
)
from mflux.models.transformer.time_text_embed import TimeTextEmbed
class Transformer(nn.Module):
def __init__(self, model_config: ModelConfig):
super().__init__()
self.pos_embed = EmbedND()
@ -26,14 +31,14 @@ class Transformer(nn.Module):
self.proj_out = nn.Linear(3072, 64)
def predict(
self,
t: int,
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,
self,
t: int,
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)
@ -51,7 +56,7 @@ class Transformer(nn.Module):
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb
rotary_embeddings=image_rotary_emb,
)
if controlnet_block_samples is not None and len(controlnet_block_samples) > 0:
interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
@ -64,7 +69,7 @@ class Transformer(nn.Module):
hidden_states = block.forward(
hidden_states=hidden_states,
text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb
rotary_embeddings=image_rotary_emb,
)
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)
@ -74,7 +79,7 @@ class Transformer(nn.Module):
+ controlnet_single_block_samples[idx // interval_control]
)
hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
hidden_states = self.proj_out(hidden_states)
noise = hidden_states

1
src/mflux/models/vae/common/attention.py
View File

@ -5,7 +5,6 @@ from mflux.config.config import Config
class Attention(nn.Module):
def __init__(self):
super().__init__()
self.group_norm = nn.GroupNorm(32, 512, pytorch_compatible=True)

25
src/mflux/models/vae/common/resnet_block_2d.py
View File

@ -5,18 +5,17 @@ from mflux.config.config import Config
class ResnetBlock2D(nn.Module):
def __init__(
self,
norm1: int,
conv1_in: int,
conv1_out: int,
norm2: int,
conv2_in: int,
conv2_out: int,
conv_shortcut_in: int | None = None,
conv_shortcut_out: int | None = None,
is_conv_shortcut: bool = False
self,
norm1: int,
conv1_in: int,
conv1_out: int,
norm2: int,
conv2_in: int,
conv2_out: int,
conv_shortcut_in: int | None = None,
conv_shortcut_out: int | None = None,
is_conv_shortcut: bool = False,
):
super().__init__()
self.norm1 = nn.GroupNorm(
@ -24,14 +23,14 @@ class ResnetBlock2D(nn.Module):
dims=norm1,
eps=1e-6,
affine=True,
pytorch_compatible=True
pytorch_compatible=True,
)
self.norm2 = nn.GroupNorm(
num_groups=32,
dims=norm2,
eps=1e-6,
affine=True,
pytorch_compatible=True
pytorch_compatible=True,
)
self.conv1 = nn.Conv2d(
in_channels=conv1_in,

3
src/mflux/models/vae/common/unet_mid_block.py
View File

@ -6,13 +6,12 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
class UnetMidBlock(nn.Module):
def __init__(self):
super().__init__()
self.attentions = [Attention()]
self.resnets = [
ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512)
ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
]
def forward(self, input_array: mx.array) -> mx.array:

1
src/mflux/models/vae/decoder/conv_in.py
View File

@ -3,7 +3,6 @@ import mlx.nn as nn
class ConvIn(nn.Module):
def __init__(self):
super().__init__()
self.conv2d = nn.Conv2d(

2
src/mflux/models/vae/decoder/conv_norm_out.py
View File

@ -12,7 +12,7 @@ class ConvNormOut(nn.Module):
dims=128,
eps=1e-6,
affine=True,
pytorch_compatible=True
pytorch_compatible=True,
)
def forward(self, input_array: mx.array) -> mx.array:

1
src/mflux/models/vae/decoder/decoder.py
View File

@ -11,7 +11,6 @@ from mflux.models.vae.decoder.up_block_4 import UpBlock4
class Decoder(nn.Module):
def __init__(self):
super().__init__()
self.conv_in = ConvIn()

3
src/mflux/models/vae/decoder/up_block_1_or_2.py
View File

@ -6,13 +6,12 @@ from mflux.models.vae.decoder.up_sampler import UpSampler
class UpBlock1Or2(nn.Module):
def __init__(self):
super().__init__()
self.resnets = [
ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512)
ResnetBlock2D(norm1=512, conv1_in=512, conv1_out=512, norm2=512, conv2_in=512, conv2_out=512),
]
self.upsamplers = [UpSampler(conv_in=512, conv_out=512)]

1
src/mflux/models/vae/decoder/up_block_3.py
View File

@ -6,7 +6,6 @@ from mflux.models.vae.decoder.up_sampler import UpSampler
class UpBlock3(nn.Module):
def __init__(self):
super().__init__()
self.resnets = [

1
src/mflux/models/vae/decoder/up_block_4.py
View File

@ -5,7 +5,6 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
class UpBlock4(nn.Module):
def __init__(self):
super().__init__()
self.resnets = [

1
src/mflux/models/vae/decoder/up_sampler.py
View File

@ -4,7 +4,6 @@ from mlx import nn
class UpSampler(nn.Module):
def __init__(self, conv_in: int, conv_out: int):
super().__init__()
self.conv = nn.Conv2d(

1
src/mflux/models/vae/encoder/conv_in.py
View File

@ -3,7 +3,6 @@ import mlx.nn as nn
class ConvIn(nn.Module):
def __init__(self):
super().__init__()
self.conv2d = nn.Conv2d(

2
src/mflux/models/vae/encoder/conv_norm_out.py
View File

@ -10,7 +10,7 @@ class ConvNormOut(nn.Module):
dims=512,
eps=1e-6,
affine=True,
pytorch_compatible=True
pytorch_compatible=True,
)
def forward(self, input_array: mx.array) -> mx.array:

1
src/mflux/models/vae/encoder/down_block_1.py
View File

@ -6,7 +6,6 @@ from mflux.models.vae.encoder.down_sampler import DownSampler
class DownBlock1(nn.Module):
def __init__(self):
super().__init__()
self.resnets = [

1
src/mflux/models/vae/encoder/down_block_2.py
View File

@ -6,7 +6,6 @@ from mflux.models.vae.encoder.down_sampler import DownSampler
class DownBlock2(nn.Module):
def __init__(self):
super().__init__()
self.resnets = [

1
src/mflux/models/vae/encoder/down_block_3.py
View File

@ -6,7 +6,6 @@ from mflux.models.vae.encoder.down_sampler import DownSampler
class DownBlock3(nn.Module):
def __init__(self):
super().__init__()
self.resnets = [

1
src/mflux/models/vae/encoder/down_block_4.py
View File

@ -5,7 +5,6 @@ from mflux.models.vae.common.resnet_block_2d import ResnetBlock2D
class DownBlock4(nn.Module):
def __init__(self):
super().__init__()
self.resnets = [

1
src/mflux/models/vae/encoder/down_sampler.py
View File

@ -4,7 +4,6 @@ from mlx import nn
class DownSampler(nn.Module):
def __init__(self, conv_in: int, conv_out: int):
super().__init__()
self.conv = nn.Conv2d(

1
src/mflux/models/vae/encoder/encoder.py
View File

@ -13,7 +13,6 @@ from mflux.models.vae.encoder.down_block_4 import DownBlock4
class Encoder(nn.Module):
def __init__(self):
super().__init__()
self.conv_in = ConvIn()

72
src/mflux/post_processing/generated_image.py
View File

@ -12,21 +12,20 @@ log = logging.getLogger(__name__)
class GeneratedImage:
def __init__(
self,
image: PIL.Image.Image,
model_config: ModelConfig,
seed: int,
prompt: str,
steps: int,
guidance: float | None,
precision: mx.Dtype,
quantization: int,
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
controlnet_strength: float | None = None,
self,
image: PIL.Image.Image,
model_config: ModelConfig,
seed: int,
prompt: str,
steps: int,
guidance: float | None,
precision: mx.Dtype,
quantization: int,
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
controlnet_strength: float | None = None,
):
self.image = image
self.model_config = model_config
@ -61,7 +60,7 @@ class GeneratedImage:
# Optionally save json metadata file
if export_json_metadata:
with open(f"{file_path.with_suffix('.json')}", 'w') as json_file:
with open(f"{file_path.with_suffix('.json')}", "w") as json_file:
json.dump(self._get_metadata(), json_file, indent=4)
# Embed metadata
@ -79,33 +78,18 @@ class GeneratedImage:
metadata_str = str(metadata)
# Convert the string to bytes (using UTF-8 encoding)
user_comment_bytes = metadata_str.encode('utf-8')
user_comment_bytes = metadata_str.encode("utf-8")
# Define the UserComment tag ID
USER_COMMENT_TAG_ID = 0x9286
# Create an EXIF dictionary
exif_dict = {
'0th': {},
'Exif': {
USER_COMMENT_TAG_ID: user_comment_bytes
},
'GPS': {},
'1st': {},
'thumbnail': None
}
# Create a piexif-compatible dictionary structure
exif_piexif_dict = {
'Exif': {
USER_COMMENT_TAG_ID: user_comment_bytes
}
}
exif_piexif_dict = {"Exif": {USER_COMMENT_TAG_ID: user_comment_bytes}}
# Load the image and embed the EXIF data
image = PIL.Image.open(path)
exif_bytes = piexif.dump(exif_piexif_dict)
image.info['exif'] = exif_bytes
image.info["exif"] = exif_bytes
# Save the image with metadata
image.save(path, exif=exif_bytes)
@ -115,15 +99,15 @@ class GeneratedImage:
def _get_metadata(self) -> dict:
return {
'model': str(self.model_config.alias),
'seed': str(self.seed),
'steps': str(self.steps),
'guidance': "None" if self.model_config == ModelConfig.FLUX1_SCHNELL else str(self.guidance),
'precision': f"{self.precision}",
'quantization': "None" if self.quantization is None else f"{self.quantization} bit",
'generation_time': f"{self.generation_time:.2f} seconds",
'lora_paths': ', '.join(self.lora_paths) if self.lora_paths else '',
'lora_scales': ', '.join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else '',
'prompt': self.prompt,
'controlnet_strength': "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
"model": str(self.model_config.alias),
"seed": str(self.seed),
"steps": str(self.steps),
"guidance": "None" if self.model_config == ModelConfig.FLUX1_SCHNELL else str(self.guidance),
"precision": f"{self.precision}",
"quantization": "None" if self.quantization is None else f"{self.quantization} bit",
"generation_time": f"{self.generation_time:.2f} seconds",
"lora_paths": ", ".join(self.lora_paths) if self.lora_paths else "",
"lora_scales": ", ".join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else "",
"prompt": self.prompt,
"controlnet_strength": "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
}

19
src/mflux/post_processing/image_util.py
View File

@ -9,17 +9,16 @@ from mflux.post_processing.generated_image import GeneratedImage
class ImageUtil:
@staticmethod
def to_image(
decoded_latents: mx.array,
seed: int,
prompt: str,
quantization: int,
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
config: RuntimeConfig,
decoded_latents: mx.array,
seed: int,
prompt: str,
quantization: int,
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
config: RuntimeConfig,
) -> GeneratedImage:
normalized = ImageUtil._denormalize(decoded_latents)
normalized_numpy = ImageUtil._to_numpy(normalized)
@ -73,7 +72,7 @@ class ImageUtil:
array = mx.transpose(array, (0, 3, 1, 2))
array = ImageUtil._normalize(array)
return array
@staticmethod
def load_image(path: str) -> Image.Image:
return Image.open(path)

10
src/mflux/save.py
View File

@ -4,10 +4,10 @@ from mflux import Flux1, ModelConfig
def main():
parser = argparse.ArgumentParser(description='Save a quantized version of Flux.1 to disk.')
parser.add_argument('--path', type=str, required=True, help='Local path for loading a model from disk')
parser.add_argument('--model', "-m", type=str, required=True, choices=["dev", "schnell"], help='The model to use ("schnell" or "dev").')
parser.add_argument('--quantize', "-q", type=int, choices=[4, 8], default=8, help='Quantize the model (4 or 8, Default is 8)')
parser = argparse.ArgumentParser(description="Save a quantized version of Flux.1 to disk.")
parser.add_argument("--path", type=str, required=True, help="Local path for loading a model from disk")
parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").")
parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=8, help="Quantize the model (4 or 8, Default is 8)")
args = parser.parse_args()
@ -23,5 +23,5 @@ def main():
print(f"Model saved at {args.path}\n")
if __name__ == '__main__':
if __name__ == "__main__":
main()

1
src/mflux/tokenizer/t5_tokenizer.py
View File

@ -3,7 +3,6 @@ from transformers import T5Tokenizer
class TokenizerT5:
def __init__(self, tokenizer: T5Tokenizer, max_length: int = 256):
self.tokenizer = tokenizer
self.max_length = max_length

17
src/mflux/tokenizer/tokenizer_handler.py
View File

@ -7,19 +7,23 @@ from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
class TokenizerHandler:
def __init__(self, repo_id: str, max_t5_length: int = 256, local_path: str | None = None):
def __init__(
self,
repo_id: str,
max_t5_length: int = 256,
local_path: str | None = None,
):
root_path = Path(local_path) if local_path else TokenizerHandler._download_or_get_cached_tokenizers(repo_id)
self.clip = transformers.CLIPTokenizer.from_pretrained(
pretrained_model_name_or_path=root_path / "tokenizer",
local_files_only=True,
max_length=TokenizerCLIP.MAX_TOKEN_LENGTH
max_length=TokenizerCLIP.MAX_TOKEN_LENGTH,
)
self.t5 = transformers.T5Tokenizer.from_pretrained(
pretrained_model_name_or_path=root_path / "tokenizer_2",
local_files_only=True,
max_length=max_t5_length
max_length=max_t5_length,
)
@staticmethod
@ -27,9 +31,6 @@ class TokenizerHandler:
return Path(
snapshot_download(
repo_id=repo_id,
allow_patterns=[
"tokenizer/**",
"tokenizer_2/**"
]
allow_patterns=["tokenizer/**", "tokenizer_2/**"],
)
)

65
src/mflux/weights/lora_converter.py
View File

@ -10,8 +10,8 @@ logger = logging.getLogger(__name__)
# This script is based on `convert_flux_lora.py` from `kohya-ss/sd-scripts`.
# For more info, see: https://github.com/kohya-ss/sd-scripts/blob/sd3/networks/convert_flux_lora.py
class LoRAConverter:
class LoRAConverter:
@staticmethod
def load_weights(lora_path: str) -> dict:
state_dict = LoRAConverter._load_pytorch_weights(lora_path)
@ -25,7 +25,6 @@ class LoRAConverter:
def _load_pytorch_weights(lora_path: str) -> dict:
state_dict = {}
with safe_open(lora_path, framework="pt") as f:
metadata = f.metadata()
for k in f.keys():
state_dict[k] = f.get_tensor(k)
return state_dict
@ -38,7 +37,7 @@ class LoRAConverter:
source,
target,
f"lora_unet_double_blocks_{i}_img_attn_proj",
f"transformer.transformer_blocks.{i}.attn.to_out.0"
f"transformer.transformer_blocks.{i}.attn.to_out.0",
)
LoRAConverter._convert_to_diffusers_cat(
source,
@ -54,25 +53,25 @@ class LoRAConverter:
source,
target,
f"lora_unet_double_blocks_{i}_img_mlp_0",
f"transformer.transformer_blocks.{i}.ff.net.0.proj"
f"transformer.transformer_blocks.{i}.ff.net.0.proj",
)
LoRAConverter._convert_to_diffusers(
source,
target,
f"lora_unet_double_blocks_{i}_img_mlp_2",
f"transformer.transformer_blocks.{i}.ff.net.2"
f"transformer.transformer_blocks.{i}.ff.net.2",
)
LoRAConverter._convert_to_diffusers(
source,
target,
f"lora_unet_double_blocks_{i}_img_mod_lin",
f"transformer.transformer_blocks.{i}.norm1.linear"
f"transformer.transformer_blocks.{i}.norm1.linear",
)
LoRAConverter._convert_to_diffusers(
source,
target,
f"lora_unet_double_blocks_{i}_txt_attn_proj",
f"transformer.transformer_blocks.{i}.attn.to_add_out"
f"transformer.transformer_blocks.{i}.attn.to_add_out",
)
LoRAConverter._convert_to_diffusers_cat(
source,
@ -88,19 +87,19 @@ class LoRAConverter:
source,
target,
f"lora_unet_double_blocks_{i}_txt_mlp_0",
f"transformer.transformer_blocks.{i}.ff_context.net.0.proj"
f"transformer.transformer_blocks.{i}.ff_context.net.0.proj",
)
LoRAConverter._convert_to_diffusers(
source,
target,
f"lora_unet_double_blocks_{i}_txt_mlp_2",
f"transformer.transformer_blocks.{i}.ff_context.net.2"
f"transformer.transformer_blocks.{i}.ff_context.net.2",
)
LoRAConverter._convert_to_diffusers(
source,
target,
f"lora_unet_double_blocks_{i}_txt_mod_lin",
f"transformer.transformer_blocks.{i}.norm1_context.linear"
f"transformer.transformer_blocks.{i}.norm1_context.linear",
)
for i in range(38):
@ -120,12 +119,13 @@ class LoRAConverter:
source,
target,
f"lora_unet_single_blocks_{i}_linear2",
f"transformer.single_transformer_blocks.{i}.proj_out"
f"transformer.single_transformer_blocks.{i}.proj_out",
)
LoRAConverter._convert_to_diffusers(
source,
target, f"lora_unet_single_blocks_{i}_modulation_lin",
f"transformer.single_transformer_blocks.{i}.norm.linear"
target,
f"lora_unet_single_blocks_{i}_modulation_lin",
f"transformer.single_transformer_blocks.{i}.norm.linear",
)
if len(source) > 0:
@ -133,12 +133,7 @@ class LoRAConverter:
return target
@staticmethod
def _convert_to_diffusers(
source: dict,
target: dict,
source_key: str,
target_key: str
):
def _convert_to_diffusers(source: dict, target: dict, source_key: str, target_key: str):
if source_key + ".lora_down.weight" not in source:
return
down_weight = source.pop(source_key + ".lora_down.weight")
@ -160,11 +155,11 @@ class LoRAConverter:
@staticmethod
def _convert_to_diffusers_cat(
source: dict,
target: dict,
source_key: str,
target_keys: list[str],
dims=None
source: dict,
target: dict,
source_key: str,
target_keys: list[str],
dims=None,
):
if source_key + ".lora_down.weight" not in source:
return
@ -204,7 +199,12 @@ class LoRAConverter:
if j == k:
continue
is_sparse = is_sparse and torch.all(
up_weight[i: i + dims[j], k * diffusers_rank: (k + 1) * diffusers_rank] == 0)
up_weight[
i : i + dims[j],
k * diffusers_rank : (k + 1) * diffusers_rank,
]
== 0
)
i += dims[j]
if is_sparse:
logger.info(f"weight is sparse: {source_key}")
@ -220,12 +220,23 @@ class LoRAConverter:
target.update({k: v for k, v in zip(diffusers_up_keys, torch.split(up_weight, dims, dim=0))})
else:
# down_weight is chunked to each split
target.update({k: v for k, v in zip(diffusers_down_keys, torch.chunk(down_weight, num_splits, dim=0))})
target.update(
{
k: v
for k, v in zip(
diffusers_down_keys,
torch.chunk(down_weight, num_splits, dim=0),
)
}
)
# up_weight is sparse: only non-zero values are copied to each split
i = 0
for j in range(len(dims)):
target[diffusers_up_keys[j]] = up_weight[i: i + dims[j], j * diffusers_rank: (j + 1) * diffusers_rank].contiguous()
target[diffusers_up_keys[j]] = up_weight[
i : i + dims[j],
j * diffusers_rank : (j + 1) * diffusers_rank,
].contiguous()
i += dims[j]
@staticmethod

26
src/mflux/weights/lora_util.py
View File

@ -6,9 +6,12 @@ log = logging.getLogger(__name__)
class LoraUtil:
@staticmethod
def apply_loras(transformer: dict, lora_files: list[str], lora_scales: list[float] | None = None) -> None:
def apply_loras(
transformer: dict,
lora_files: list[str],
lora_scales: list[float] | None = None,
) -> None:
lora_scales = LoraUtil._validate_lora_scales(lora_files, lora_scales)
for lora_file, lora_scale in zip(lora_files, lora_scales):
@ -29,6 +32,7 @@ class LoraUtil:
@staticmethod
def _apply_lora(transformer: dict, lora_file: str, lora_scale: float) -> None:
from mflux.weights.weight_handler import WeightHandler
lora_transformer, _ = WeightHandler.load_transformer(lora_path=lora_file)
LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale)
@ -54,14 +58,14 @@ class LoraUtil:
visiting.append(splitKey)
else:
parentKey = ".".join(visiting)
if parentKey in visited and 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]:
if parentKey in visited and "lora_A" in visited[parentKey] and "lora_B" in visited[parentKey]:
continue
if not splitKey.startswith("lora_"):
visiting.append(splitKey)
parentKey = ".".join(visiting)
if splitKey == "net":
target['net'] = list({})
target = target['net']
target["net"] = list({})
target = target["net"]
elif splitKey == "0":
target.append({})
target = target[0]
@ -74,11 +78,11 @@ class LoraUtil:
if parentKey not in visited:
visited[parentKey] = {}
visited[parentKey][splitKey] = weight
if not 'weight' in target:
if "weight" not in target:
raise ValueError(f"LoRA weights for layer {parentKey} cannot be loaded into the model.")
if 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]:
lora_a = visited[parentKey]['lora_A']
lora_b = visited[parentKey]['lora_B']
transWeight = target['weight']
if "lora_A" in visited[parentKey] and "lora_B" in visited[parentKey]:
lora_a = visited[parentKey]["lora_A"]
lora_b = visited[parentKey]["lora_B"]
transWeight = target["weight"]
weight = transWeight + lora_scale * (lora_b @ lora_a)
target['weight'] = weight
target["weight"] = weight

7
src/mflux/weights/model_saver.py
View File

@ -7,7 +7,6 @@ from transformers import CLIPTokenizer, T5Tokenizer
class ModelSaver:
@staticmethod
def save_model(model, bits: int, base_path: str):
# Save the tokenizers
@ -32,7 +31,11 @@ class ModelSaver:
path.mkdir(parents=True, exist_ok=True)
weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters())))
for i, weight in enumerate(weights):
mx.save_safetensors(str(path / f"{i}.safetensors"), weight, {"quantization_level": str(bits)})
mx.save_safetensors(
str(path / f"{i}.safetensors"),
weight,
{"quantization_level": str(bits)},
)
@staticmethod
def _split_weights(base_path: str, weights: dict, max_file_size_gb: int = 2) -> list:

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

@ -11,13 +11,12 @@ from mflux.weights.weight_util import WeightUtil
class WeightHandler:
def __init__(
self,
repo_id: str | None = None,
local_path: str | None = None,
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
self,
repo_id: str | None = None,
local_path: str | None = None,
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
):
root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
@ -66,7 +65,7 @@ class WeightHandler:
weights, quantization_level = WeightHandler._get_weights("transformer", root_path, lora_path)
if lora_path:
if 'transformer' not in weights:
if "transformer" not in weights:
weights = LoRAConverter.load_weights(lora_path)
weights = weights["transformer"]
@ -80,18 +79,23 @@ class WeightHandler:
if block.get("ff") is not None:
block["ff"] = {
"linear1": block["ff"]["net"][0]["proj"],
"linear2": block["ff"]["net"][2]
"linear2": block["ff"]["net"][2],
}
if block.get("ff_context") is not None:
block["ff_context"] = {
"linear1": block["ff_context"]["net"][0]["proj"],
"linear2": block["ff_context"]["net"][2]
"linear2": block["ff_context"]["net"][2],
}
return weights, quantization_level
@staticmethod
def load_controlnet_transformer(controlnet_id: str) -> (dict, int):
controlnet_path = Path(snapshot_download(repo_id=controlnet_id,allow_patterns=["*.safetensors","config.json"]))
controlnet_path = Path(
snapshot_download(
repo_id=controlnet_id,
allow_patterns=["*.safetensors", "config.json"],
)
)
file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors"))
quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
weights = list(mx.load(str(file)).items())
@ -112,12 +116,12 @@ class WeightHandler:
for block in weights["transformer_blocks"]:
block["ff"] = {
"linear1": block["ff"]["net"][0]["proj"],
"linear2": block["ff"]["net"][2]
"linear2": block["ff"]["net"][2],
}
if block.get("ff_context") is not None:
block["ff_context"] = {
"linear1": block["ff_context"]["net"][0]["proj"],
"linear2": block["ff_context"]["net"][2]
"linear2": block["ff_context"]["net"][2],
}
config = json.load(open(controlnet_path / "config.json"))
return weights, quantization_level, config
@ -131,16 +135,20 @@ class WeightHandler:
return weights, quantization_level
# Reshape and process the huggingface weights
weights['decoder']['conv_in'] = {'conv2d': weights['decoder']['conv_in']}
weights['decoder']['conv_out'] = {'conv2d': weights['decoder']['conv_out']}
weights['decoder']['conv_norm_out'] = {'norm': weights['decoder']['conv_norm_out']}
weights['encoder']['conv_in'] = {'conv2d': weights['encoder']['conv_in']}
weights['encoder']['conv_out'] = {'conv2d': weights['encoder']['conv_out']}
weights['encoder']['conv_norm_out'] = {'norm': weights['encoder']['conv_norm_out']}
weights["decoder"]["conv_in"] = {"conv2d": weights["decoder"]["conv_in"]}
weights["decoder"]["conv_out"] = {"conv2d": weights["decoder"]["conv_out"]}
weights["decoder"]["conv_norm_out"] = {"norm": weights["decoder"]["conv_norm_out"]}
weights["encoder"]["conv_in"] = {"conv2d": weights["encoder"]["conv_in"]}
weights["encoder"]["conv_out"] = {"conv2d": weights["encoder"]["conv_out"]}
weights["encoder"]["conv_norm_out"] = {"norm": weights["encoder"]["conv_norm_out"]}
return weights, quantization_level
@staticmethod
def _get_weights(model_name: str, root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
def _get_weights(
model_name: str,
root_path: Path | None = None,
lora_path: str | None = None,
) -> (dict, int):
weights = []
quantization_level = None
@ -174,6 +182,6 @@ class WeightHandler:
"text_encoder_2/*.safetensors",
"transformer/*.safetensors",
"vae/*.safetensors",
]
],
)
)

1
src/mflux/weights/weight_util.py
View File

@ -2,7 +2,6 @@ from mflux.config.config import Config
class WeightUtil:
@staticmethod
def flatten(params):
return [(k, v) for p in params for (k, v) in p]