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

Rewrite weight handling and add quantization support

Browse Source
This commit is contained in:
filipstrand 2024-08-30 14:04:58 +02:00
parent 478fc4f8bc
commit fdc4cb9892
14 changed files with 335 additions and 141 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

132
README.md
View File

@ -37,19 +37,19 @@ like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) f
```
### Generating an image
Run the provided [main.py](main.py) by specifying a prompt and some optional arguments like so using the default `Schnell` model:
Run the provided [main.py](main.py) by specifying a prompt and some optional arguments like so using the `schnell` model:
```
python main.py --prompt "Luxury food photograph" --steps 2 --seed 2
python main.py --model schnell --prompt "Luxury food photograph" --steps 2 --seed 2
```
or use the slower, but more powerful `Dev` model and run it with more time steps:
or use the slower, but more powerful `dev` model and run it with more time steps:
```
python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
```
⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively).* ⚠️
⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#quantization) section for running compressed versions of the model.* ⚠️
*By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:*
@ -65,11 +65,11 @@ python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
- **`--prompt`** (required, `str`): Text description of the image to generate.
- **`--model`** or **`-m`** (required, `str`): Model to use for generation (`"schnell"` or `"dev"`).
- **`--output`** (optional, `str`, default: `"image.png"`): Output image filename.
- **`--model`** (optional, `str`, default: `"schnell"`): Model to use for generation (`"schnell"` or `"dev"`).
- **`--seed`** (optional, `int`, default: `0`): Seed for random number generation. Default is time-based.
- **`--seed`** (optional, `int`, default: `None`): Seed for random number generation. Default is time-based.
- **`--height`** (optional, `int`, default: `1024`): Height of the output image in pixels.
@ -79,6 +79,11 @@ python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
- **`--guidance`** (optional, `float`, default: `3.5`): Guidance scale (only used for `"dev"` model).
- **`--path`** (optional, `str`, default: `None`): Path to a local model on disk.
- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#quantization) (choose between `4` or `8`).
Or make a new separate script like the following
```python
@ -90,7 +95,7 @@ from flux_1.config.config import Config
from flux_1.flux import Flux1
from flux_1.post_processing.image_util import ImageUtil
flux = Flux1.from_alias("schnell") # "schnell" or "dev"
flux = Flux1.from_alias(alias="schnell") # "schnell" or "dev"
image = flux.generate_image(
seed=3,
@ -113,6 +118,7 @@ To time your machine, run the following:
```
time python main.py \
--prompt "Luxury food photograph" \
--model schnell \
--steps 2 \
--seed 2 \
--height 1024 \
@ -180,7 +186,114 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
---
### Quantization
MFLUX supports running FLUX in 4-bit or 8-bit quantized mode. Running a quantized version can greatly speed up the
generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#size-comparisons-for-quantized-models).
```
python main.py \
--model schnell \
--steps 2 \
--seed 2 \
--quantize 8 \
--height 1920 \
--width 1024 \
--prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
```
![image](src/flux_1/assets/comparison6.jpg)
*In this example, weights are quantized at **runtime** - this is convenient if you don't want to [save a quantized copy of the weights to disk](#saving-a-quantized-version-to-disk), but still want to benefit from the potential speedup and RAM reduction quantization might bring.*
By selecting the `--quantize` or `-q` flag to be `4`, `8`, or removing it entirely, we get all 3 images above. As can be seen, there is very little difference between the images (especially between the 8-bit, and the non-quantized result).
Image generation times in this example are based on a 2021 M1 Pro (32GB) machine. Even though the images are almost identical, there is a ~2x speedup by
running the 8-bit quantized version on this particular machine. Unlike the non-quantized version, for the 8-bit version the swap memory usage is drastically reduced and GPU utilization is close to 100% during the whole generation. Results here can vary across different machines.
#### Size comparisons for quantized models
The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
| 4 bit | 8 bit | Original (16 bit) |
|--------|---------|-------------------|
| 9.85GB | 18.16GB | 33.73GB |
The reason weights sizes are not fully cut in half is because a small number of weights are not quantized and kept at full precision.
#### Saving a quantized version to disk
To save a local copy of the quantized weights, run the `save.py` script like so:
```
python save.py \
--path "/Users/filipstrand/Desktop/schnell_8bit" \
--model schnell \
--quantize 8
```
*Note that when saving a quantized version, you will need the original huggingface weights.*
#### Loading and running a quantized version from disk
To generate a new image from the quantized model, simply provide a `--path` to where it was saved:
```
python main.py \
--path "/Users/filipstrand/Desktop/schnell_8bit" \
--model schnell \
--steps 2 \
--seed 2 \
--height 1920 \
--width 1024 \
--prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
```
*Note: Once we have a local model (quantized or not) specified via the `--path` argument, the huggingface cache models are not required to launch the model*.
### Running a non-quantized model directly from disk
MFLUX also supports running a non-quantized model directly from a custom location.
In the example below, the model is placed in `/Users/filipstrand/Desktop/schnell`:
```
python main.py \
--path "/Users/filipstrand/Desktop/schnell" \
--model schnell \
--steps 2 \
--seed 2 \
--prompt "Luxury food photograph"
```
Note that the `--model` flag must be set when loading a model from disk.
Also note that unlike when using the typical `alias` way of initializing the model (which internally handles that the required resources are downloaded),
when loading a model directly from disk, we require the downloaded models to look like the following:
```
.
├── text_encoder
│   └── model.safetensors
├── text_encoder_2
│   ├── model-00001-of-00002.safetensors
│   └── model-00002-of-00002.safetensors
├── tokenizer
│   ├── merges.txt
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   └── vocab.json
├── tokenizer_2
│   ├── special_tokens_map.json
│   ├── spiece.model
│   ├── tokenizer.json
│   └── tokenizer_config.json
├── transformer
│   ├── diffusion_pytorch_model-00001-of-00003.safetensors
│   ├── diffusion_pytorch_model-00002-of-00003.safetensors
│   └── diffusion_pytorch_model-00003-of-00003.safetensors
└── vae
└── diffusion_pytorch_model.safetensors
```
This mirrors how the resources are placed in the [HuggingFace Repo](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main) for FLUX.1.
### Current limitations
@ -191,5 +304,4 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
- [ ] LoRA adapters
- [ ] LoRA fine-tuning
- [ ] Frontend support (Gradio/Streamlit/Other?)
- [ ] Support for quantized models
- [ ] Frontend support (Gradio/Streamlit/Other?)

14
main.py
View File

@ -5,6 +5,7 @@ import time
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
from flux_1.config.model_config import ModelConfig
from flux_1.config.config import Config
from flux_1.flux import Flux1
from flux_1.post_processing.image_util import ImageUtil
@ -14,18 +15,27 @@ 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', type=str, default="schnell", help='The model to use ("schnell" or "dev"). Default is "schnell".')
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=4, 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')
args = parser.parse_args()
if args.path and args.model is None:
parser.error("--model must be specified when using --path")
seed = int(time.time()) if args.seed is None else args.seed
flux = Flux1.from_alias(args.model)
flux = Flux1(
model_config=ModelConfig.from_alias(args.model),
quantize_full_weights=args.quantize,
local_path=args.path
)
image = flux.generate_image(
seed=seed,

3
requirements.txt
View File

@ -5,4 +5,5 @@ transformers>=4.44.0
sentencepiece>=0.2.0
torch>=2.3.1
tqdm>=4.66.5
huggingface-hub>=0.24.5
huggingface-hub>=0.24.5
safetensors>=0.4.4

32
save.py Normal file
View File

@ -0,0 +1,32 @@
import argparse
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
from flux_1.flux import Flux1
from flux_1.config.model_config import 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)')
args = parser.parse_args()
print(f"Saving model {args.model} with quantization level {args.quantize}\n")
flux = Flux1(
model_config=ModelConfig.from_alias(args.model),
quantize_full_weights=args.quantize,
)
flux.save_model(args.path)
print(f"Model saved at {args.path}\n")
if __name__ == '__main__':
main()

BIN
src/flux_1/assets/comparison6.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.0 MiB

9
src/flux_1/config/model_config.py
View File

@ -17,15 +17,6 @@ class ModelConfig(Enum):
self.num_train_steps = num_train_steps
self.max_sequence_length = max_sequence_length
@staticmethod
def from_repo(model_name: str) -> "ModelConfig":
try:
for model in ModelConfig:
if model.model_name == model_name:
return model
except KeyError:
raise ValueError(f"'{model_name}' is not a valid model")
@staticmethod
def from_alias(alias: str) -> "ModelConfig":
try:

113
src/flux_1/flux.py
View File

@ -1,6 +1,10 @@
from pathlib import Path
import PIL
import mlx.core as mx
from PIL import Image
from mlx import nn
from mlx.utils import tree_flatten
from tqdm import tqdm
from flux_1.config.config import Config
@ -19,47 +23,63 @@ from flux_1.weights.weight_handler import WeightHandler
class Flux1:
def __init__(self, repo_id: str):
self.model_config = ModelConfig.from_repo(repo_id)
def __init__(
self,
model_config: ModelConfig,
quantize_full_weights: int | None = None,
local_path: str | None = None,
):
self.model_config = model_config
self.quantize_full_weights = quantize_full_weights
# Initialize the tokenizers
tokenizers = TokenizerHandler.load_from_disk_or_huggingface(repo_id, self.model_config.max_sequence_length)
# Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
# Initialize the models
weights = WeightHandler.load_from_disk_or_huggingface(repo_id)
self.vae = VAE(weights.vae)
self.transformer = Transformer(weights.transformer)
self.t5_text_encoder = T5Encoder(weights.t5_encoder)
self.clip_text_encoder = CLIPEncoder(weights.clip_encoder)
self.vae = VAE()
self.transformer = Transformer(model_config)
self.t5_text_encoder = T5Encoder()
self.clip_text_encoder = CLIPEncoder()
@staticmethod
def from_repo(repo_id: str) -> "Flux1":
return Flux1(repo_id)
# Load the weights from disk, huggingface cache, or download from huggingface
weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path)
@staticmethod
def from_alias(alias: str) -> "Flux1":
return Flux1(ModelConfig.from_alias(alias).model_name)
# Set the loaded weights if they are not quantized
if weights.quantization_level is None:
self._set_model_weights(weights)
# Optionally quantize the model here at initialization (also required if about to load quantized weights)
if quantize_full_weights is not None or weights.quantization_level is not None:
bits = weights.quantization_level if weights.quantization_level is not None else quantize_full_weights
nn.quantize(self.vae, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
nn.quantize(self.transformer, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 64, group_size=64, bits=bits)
nn.quantize(self.t5_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
nn.quantize(self.clip_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
# If loading previously saved quantized weights, the weights must be set after modules have been quantized
if weights.quantization_level is not None:
self._set_model_weights(weights)
def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image:
# Create a new runtime config based on the model type and input parameters
config = RuntimeConfig(config, self.model_config)
# Create the latents
# 1. Create the initial latents
latents = mx.random.normal(
shape=[1, (config.height // 16) * (config.width // 16), 64],
key=mx.random.key(seed)
)
# Embedd the prompt
# 2. Embedd the prompt
t5_tokens = self.t5_tokenizer.tokenize(prompt)
clip_tokens = self.clip_tokenizer.tokenize(prompt)
prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
for t in tqdm(range(config.num_inference_steps)):
# Predict the noise
# 3.t Predict the noise
noise = self.transformer.predict(
t=t,
prompt_embeds=prompt_embeds,
@ -68,14 +88,14 @@ class Flux1:
config=config,
)
# Take one denoise step
# 4.t Take one denoise step
dt = config.sigmas[t + 1] - config.sigmas[t]
latents += noise * dt
# To enable progress tracking
# Evaluate to enable progress tracking
mx.eval(latents)
# Decode the latent array
# 5. Decode the latent array
latents = Flux1._unpack_latents(latents, config.height, config.width)
decoded = self.vae.decode(latents)
return ImageUtil.to_image(decoded)
@ -87,10 +107,49 @@ class Flux1:
latents = mx.reshape(latents, (1, 16, width // 16 * 2, height // 16 * 2))
return latents
def encode(self, path: str) -> mx.array:
array = ImageUtil.to_array(Image.open(path))
return self.vae.encode(array)
@staticmethod
def from_alias(alias: str) -> "Flux1":
return Flux1(ModelConfig.from_alias(alias))
def decode(self, code: mx.array) -> PIL.Image.Image:
decoded = self.vae.decode(code)
return ImageUtil.to_image(decoded)
def _set_model_weights(self, weights):
self.vae.update(weights.vae)
self.transformer.update(weights.transformer)
self.t5_text_encoder.update(weights.t5_encoder)
self.clip_text_encoder.update(weights.clip_encoder)
def save_model(self, base_path: str):
def _save_tokenizer(tokenizer, subdir: str):
path = Path(base_path) / subdir
path.mkdir(parents=True, exist_ok=True)
tokenizer.save_pretrained(path)
def _save_weights(model, subdir: str):
path = Path(base_path) / subdir
path.mkdir(parents=True, exist_ok=True)
weights = _split_weights(dict(tree_flatten(model.parameters())))
for i, weight in enumerate(weights):
mx.save_safetensors(str(path / f"{i}.safetensors"), weight, {"quantization_level": str(self.quantize_full_weights)})
def _split_weights(weights: dict, max_file_size_gb: int = 2) -> list:
# Copied from mlx-examples repo
max_file_size_bytes = max_file_size_gb << 30
shards = []
shard, shard_size = {}, 0
for k, v in weights.items():
if shard_size + v.nbytes > max_file_size_bytes:
shards.append(shard)
shard, shard_size = {}, 0
shard[k] = v
shard_size += v.nbytes
shards.append(shard)
return shards
# Save the tokenizers
_save_tokenizer(self.clip_tokenizer.tokenizer, "tokenizer")
_save_tokenizer(self.t5_tokenizer.tokenizer, "tokenizer_2")
# Save the models
_save_weights(self.vae, "vae")
_save_weights(self.transformer, "transformer")
_save_weights(self.clip_text_encoder, "text_encoder")
_save_weights(self.t5_text_encoder, "text_encoder_2")

5
src/flux_1/models/text_encoder/clip_encoder/clip_encoder.py
View File

@ -6,13 +6,10 @@ from flux_1.models.text_encoder.clip_encoder.clip_text_model import CLIPTextMode
class CLIPEncoder(nn.Module):
def __init__(self, weights: dict):
def __init__(self):
super().__init__()
self.text_model = CLIPTextModel(dims=768, num_encoder_layers=12)
# Load the weights after all components are initialized
self.update(weights)
def forward(self, tokens: mx.array) -> mx.array:
pooled_output = self.text_model.forward(tokens)
return pooled_output

5
src/flux_1/models/text_encoder/t5_encoder/t5_encoder.py
View File

@ -7,15 +7,12 @@ from flux_1.models.text_encoder.t5_encoder.t5_layer_norm import T5LayerNorm
class T5Encoder(nn.Module):
def __init__(self, weights: dict):
def __init__(self):
super().__init__()
self.shared = nn.Embedding(num_embeddings=32128, dims=4096)
self.t5_blocks = [T5Block(i) for i in range(24)]
self.final_layer_norm = T5LayerNorm()
# Load the weights after all components are initialized
self.update(weights)
def forward(self, tokens: mx.array):
hidden_states = self.shared(tokens)
for block in self.t5_blocks:

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

@ -3,6 +3,7 @@ from mlx import nn
import mlx.core as mx
from flux_1.config.config import Config
from flux_1.config.model_config import ModelConfig
from flux_1.models.transformer.text_embedder import TextEmbedder
from flux_1.models.transformer.timestep_embedder import TimestepEmbedder
from flux_1.models.transformer.guidance_embedder import GuidanceEmbedder
@ -10,18 +11,16 @@ from flux_1.models.transformer.guidance_embedder import GuidanceEmbedder
class TimeTextEmbed(nn.Module):
def __init__(self, with_guidance_embed: bool = False):
def __init__(self, model_config: ModelConfig):
super().__init__()
self.text_embedder = TextEmbedder()
self.with_guidance_embed = with_guidance_embed
if self.with_guidance_embed:
self.guidance_embedder = GuidanceEmbedder()
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:
time_steps_proj = self._time_proj(time_step)
time_steps_emb = self.timestep_embedder.forward(time_steps_proj)
if self.with_guidance_embed:
if self.guidance_embedder is not None:
time_steps_emb += self.guidance_embedder.forward(self._time_proj(guidance))
pooled_projections = self.text_embedder.forward(pooled_projection)
conditioning = time_steps_emb + pooled_projections

10
src/flux_1/models/transformer/transformer.py
View File

@ -1,7 +1,7 @@
import mlx.core as mx
from mlx import nn
from flux_1.config.config import Config
from flux_1.config.model_config import ModelConfig
from flux_1.config.runtime_config import RuntimeConfig
from flux_1.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
from flux_1.models.transformer.embed_nd import EmbedND
@ -12,21 +12,17 @@ from flux_1.models.transformer.time_text_embed import TimeTextEmbed
class Transformer(nn.Module):
def __init__(self, weights: dict):
def __init__(self, model_config: ModelConfig):
super().__init__()
self.pos_embed = EmbedND()
self.x_embedder = nn.Linear(64, 3072)
with_guidance_embed = "guidance_embedder" in weights["time_text_embed"].keys()
self.time_text_embed = TimeTextEmbed(with_guidance_embed=with_guidance_embed)
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(19)]
self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)]
self.norm_out = AdaLayerNormContinuous(3072, 3072)
self.proj_out = nn.Linear(3072, 64)
# Load the weights after all components are initialized
self.update(weights)
def predict(
self,
t: int,

5
src/flux_1/models/vae/vae.py
View File

@ -9,14 +9,11 @@ class VAE(nn.Module):
scaling_factor: int = 0.3611
shift_factor: int = 0.1159
def __init__(self, weights: dict):
def __init__(self):
super().__init__()
self.decoder = Decoder()
self.encoder = Encoder()
# Load the weights after all components are initialized
self.update(weights)
def decode(self, latents: mx.array) -> mx.array:
scaled_latents = (latents / self.scaling_factor) + self.shift_factor
return self.decoder.decode(scaled_latents)

9
src/flux_1/tokenizer/tokenizer_handler.py
View File

@ -4,13 +4,12 @@ import transformers
from huggingface_hub import snapshot_download
from flux_1.tokenizer.clip_tokenizer import TokenizerCLIP
from flux_1.tokenizer.t5_tokenizer import TokenizerT5
class TokenizerHandler:
def __init__(self, repo_id: str, max_t5_length: int = 256):
root_path = TokenizerHandler._download_or_get_cached_tokenizers(repo_id)
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",
@ -23,10 +22,6 @@ class TokenizerHandler:
max_length=max_t5_length
)
@staticmethod
def load_from_disk_or_huggingface(repo_id: str, max_t5_length: int = 256) -> "TokenizerHandler":
return TokenizerHandler(repo_id, max_t5_length)
@staticmethod
def _download_or_get_cached_tokenizers(repo_id: str) -> Path:
return Path(

130
src/flux_1/weights/weight_handler.py
View File

@ -3,83 +3,67 @@ from pathlib import Path
import mlx.core as mx
from huggingface_hub import snapshot_download
from mlx.utils import tree_unflatten
from safetensors import safe_open
from flux_1.config.config import Config
class WeightHandler:
def __init__(self, repo_id: str):
root_path = WeightHandler._download_or_get_cached_weights(repo_id)
def __init__(
self,
repo_id: str | None = None,
local_path: str | None = None,
):
root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
self.clip_encoder = WeightHandler._clip_encoder(
weights=WeightHandler._load(root_path / "text_encoder/model.safetensors")
)
self.t5_encoder = WeightHandler._t5_encoder(
weights_1=WeightHandler._load(root_path / "text_encoder_2/model-00001-of-00002.safetensors"),
weights_2=WeightHandler._load(root_path / "text_encoder_2/model-00002-of-00002.safetensors"),
)
self.vae = WeightHandler._vae(
weights=WeightHandler._load(root_path / "vae/diffusion_pytorch_model.safetensors")
)
self.transformer = WeightHandler._transformer(
weights_1=WeightHandler._load(root_path / "transformer/diffusion_pytorch_model-00001-of-00003.safetensors"),
weights_2=WeightHandler._load(root_path / "transformer/diffusion_pytorch_model-00002-of-00003.safetensors"),
weights_3=WeightHandler._load(root_path / "transformer/diffusion_pytorch_model-00003-of-00003.safetensors"),
)
self.clip_encoder, _ = WeightHandler._clip_encoder(root_path=root_path)
self.t5_encoder, _ = WeightHandler._t5_encoder(root_path=root_path)
self.vae, _ = WeightHandler._vae(root_path=root_path)
self.transformer, self.quantization_level = WeightHandler._transformer(root_path=root_path)
@staticmethod
def load_from_disk_or_huggingface(repo_id: str) -> "WeightHandler":
return WeightHandler(repo_id)
def _clip_encoder(root_path: Path) -> (dict, int):
weights, quantization_level = WeightHandler._get_weights("text_encoder", root_path)
return weights, quantization_level
@staticmethod
def _load(path: Path) -> list[dict]:
return list(mx.load(str(path)).items())
def _t5_encoder(root_path: Path) -> (dict, int):
weights, quantization_level = WeightHandler._get_weights("text_encoder_2", root_path)
@staticmethod
def _clip_encoder(weights: list[dict]) -> dict:
weights = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights])
unflatten = tree_unflatten(weights)
return unflatten
# Quantized weights (i.e. ones exported from this project) don't need any post-processing.
if quantization_level is not None:
return weights, quantization_level
@staticmethod
def _t5_encoder(
weights_1: list[dict],
weights_2: list[dict],
) -> dict:
weights_1 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_1])
weights_2 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_2])
unflatten = tree_unflatten(weights_1 + weights_2)
unflatten["final_layer_norm"] = unflatten["encoder"]["final_layer_norm"]
for block in unflatten["encoder"]["block"]:
# Reshape and process the huggingface weights
weights["final_layer_norm"] = weights["encoder"]["final_layer_norm"]
for block in weights["encoder"]["block"]:
attention = block["layer"][0]
ff = block["layer"][1]
block.pop("layer")
block["attention"] = attention
block["ff"] = ff
unflatten["t5_blocks"] = unflatten["encoder"]["block"]
weights["t5_blocks"] = weights["encoder"]["block"]
# Only the first layer has the weights for "relative_attention_bias", we duplicate them here to keep code simple
relative_attention_bias = unflatten["t5_blocks"][0]["attention"]["SelfAttention"]["relative_attention_bias"]
for block in unflatten["t5_blocks"][1:]:
relative_attention_bias = weights["t5_blocks"][0]["attention"]["SelfAttention"]["relative_attention_bias"]
for block in weights["t5_blocks"][1:]:
block["attention"]["SelfAttention"]["relative_attention_bias"] = relative_attention_bias
unflatten.pop("encoder")
return unflatten
weights.pop("encoder")
return weights, quantization_level
@staticmethod
def _transformer(
weights_1: list[dict],
weights_2: list[dict],
weights_3: list[dict],
) -> dict:
weights_1 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_1])
weights_2 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_2])
weights_3 = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights_3])
unflatten = tree_unflatten(weights_1 + weights_2 + weights_3)
def _transformer(root_path: Path) -> (dict, int):
weights, quantization_level = WeightHandler._get_weights("transformer", root_path)
for block in unflatten["transformer_blocks"]:
# Quantized weights (i.e. ones exported from this project) don't need any post-processing.
if quantization_level is not None:
return weights, quantization_level
# Reshape and process the huggingface weights
for block in weights["transformer_blocks"]:
block["ff"] = {
"linear1": block["ff"]["net"][0]["proj"],
"linear2": block["ff"]["net"][2]
@ -89,19 +73,43 @@ class WeightHandler:
"linear1": block["ff_context"]["net"][0]["proj"],
"linear2": block["ff_context"]["net"][2]
}
return unflatten
return weights, quantization_level
@staticmethod
def _vae(weights: list[dict]) -> dict:
weights = WeightHandler._flatten([WeightHandler._reshape_weights(k, v) for k, v in weights])
def _vae(root_path: Path) -> (dict, int):
weights, quantization_level = WeightHandler._get_weights("vae", root_path)
# Quantized weights (i.e. ones exported from this project) don't need any post-processing.
if quantization_level is not None:
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']}
return weights, quantization_level
@staticmethod
def _get_weights(model_name: str, root_path: Path) -> (dict, int):
weights = []
quantization_level = None
for file in sorted(root_path.glob(model_name + "/*.safetensors")):
quantization_level = safe_open(file, framework="pt").metadata().get("quantization_level")
weight = list(mx.load(str(file)).items())
weights.extend(weight)
# Non huggingface weights (i.e. ones exported from this project) don't need any reshaping.
if quantization_level is not None:
return tree_unflatten(weights), quantization_level
# Huggingface weights needs to be reshaped
weights = [WeightHandler._reshape_weights(k, v) for k, v in weights]
weights = WeightHandler._flatten(weights)
unflatten = tree_unflatten(weights)
unflatten['decoder']['conv_in'] = {'conv2d': unflatten['decoder']['conv_in']}
unflatten['decoder']['conv_out'] = {'conv2d': unflatten['decoder']['conv_out']}
unflatten['decoder']['conv_norm_out'] = {'norm': unflatten['decoder']['conv_norm_out']}
unflatten['encoder']['conv_in'] = {'conv2d': unflatten['encoder']['conv_in']}
unflatten['encoder']['conv_out'] = {'conv2d': unflatten['encoder']['conv_out']}
unflatten['encoder']['conv_norm_out'] = {'norm': unflatten['encoder']['conv_norm_out']}
return unflatten
return unflatten, quantization_level
@staticmethod
def _flatten(params):