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

Add proper image class and save metadata

Browse Source
This commit is contained in:
filipstrand 2024-09-06 19:19:05 +02:00
parent 97dc64436e
commit fab10cc046
9 changed files with 210 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -12,3 +12,4 @@
*.jpg
*.pyc
*.safetensors
*.json

40
README.md
View File

@ -8,10 +8,10 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
### Philosophy
MFLUX (MacFLUX) is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx).
MFLUX is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx).
MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used
except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models
(thereby avoiding too many abstractions). MFLUX priorities readability over generality and performance.
(thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#image-generation-speed-updated), [and even faster quantized](#quantization).
All models are implemented from scratch in MLX and only the tokenizers are used via the
[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies
@ -40,13 +40,13 @@ like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) f
Run the provided [main.py](main.py) by specifying a prompt and some optional arguments like so using the `schnell` model:
```
python main.py --model schnell --prompt "Luxury food photograph" --steps 2 --seed 2
python main.py --model schnell --prompt "Luxury food photograph" --steps 2 --seed 2 -q 8
```
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
python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 -q 8
```
⚠️ *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.* ⚠️
@ -86,8 +86,10 @@ python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
- **`--lora-scales`** (optional, `[float]`, default: `None`): The scale for each respective [LoRA](#LoRA) (will default to `1.0` if not specified and only one LoRA weight is loaded.)
- **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `exiftool image.png`)
Or make a new separate script like the following
Or, with the correct python environment active, make a new separate script like the following:
```python
import sys
@ -96,27 +98,29 @@ sys.path.append("/path/to/mflux/src")
from flux_1.config.config import Config
from flux_1.flux import Flux1
from flux_1.post_processing.image_util import ImageUtil
# Load the model
flux = Flux1.from_alias(alias="schnell") # "schnell" or "dev"
# Generate an image
image = flux.generate_image(
seed=3,
prompt="Luxury food photograph of a birthday cake. In the middle it has three candles shaped like letters spelling the word 'MLX'. It has perfect lighting and a cozy background with big bokeh and shallow depth of field. The mood is a sunset balcony in tuscany. The photo is taken from the side of the cake. The scene is complemented by a warm, inviting light that highlights the textures and colors of the ingredients, giving it an appetizing and elegant look.",
seed=2,
prompt="Luxury food photograph",
config=Config(
num_inference_steps=2, # Schnell works well with 2-4 steps, Dev works well with 20-25 steps
height=768,
width=1360,
num_inference_steps=2, # "schnell" works well with 2-4 steps, "dev" works well with 20-25 steps
height=1024,
width=1024,
)
)
ImageUtil.save_image(image, "image.png")
image.save(path="image.png")
```
For more options on how to configure MFLUX, please see [main.py](main.py).
### Image generation speed (updated)
These numbers are based on the Schnell model, with the configuration provided in the code snippet below.
These numbers are based on the non-quantized `schnell` model, with the configuration provided in the code snippet below.
To time your machine, run the following:
```
time python main.py \
@ -140,6 +144,9 @@ time python main.py \
| 2021 M1 Pro (32GB) | @filipstrand | ~160s | |
| 2023 M2 Max (32GB) | @filipstrand | ~70s | |
*Note that these numbers includes starting the application from scratch, which means doing model i/o, setting/quantizing weights etc.
If we assume that the model is already loaded, you can inspect the image metadata using `exiftool image.png` and see the total duration of the denoising loop (excluding text embedding).*
### Equivalent to Diffusers implementation
There is only a single source of randomness when generating an image: The initial latent array.
@ -251,7 +258,9 @@ python main.py \
--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*.
*Note: When loading a quantized model from disk, there is no need to pass in `-q` flag, since we can infer this from the weight metadata.*
*Also Note: Once we have a local model (quantized [or not](#running-a-non-quantized-model-directly-from-disk)) specified via the `--path` argument, the huggingface cache models are not required to launch the model*.
### Running a non-quantized model directly from disk
@ -297,6 +306,9 @@ when loading a model directly from disk, we require the downloaded models to loo
└── 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.
*Huggingface weights, unlike quantized ones exported directly from this project, have to be
processed a bit differently, which is why we require this structure above.*
### LoRA

9
main.py
View File

@ -8,7 +8,6 @@ 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
def main():
@ -25,20 +24,23 @@ def main():
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()
if args.path and args.model is None:
parser.error("--model must be specified when using --path")
# Load the model
flux = Flux1(
model_config=ModelConfig.from_alias(args.model),
quantize_full_weights=args.quantize,
quantize=args.quantize,
local_path=args.path,
lora_paths=args.lora_paths,
lora_scales=args.lora_scales
)
# Generate an image
image = flux.generate_image(
seed=int(time.time()) if args.seed is None else args.seed,
prompt=args.prompt,
@ -50,7 +52,8 @@ def main():
)
)
ImageUtil.save_image(image, args.output)
# Save the image
image.save(path=args.output, export_json_metadata=args.metadata)
if __name__ == '__main__':

3
requirements.txt
View File

@ -6,4 +6,5 @@ sentencepiece>=0.2.0
torch>=2.3.1
tqdm>=4.66.5
huggingface-hub>=0.24.5
safetensors>=0.4.4
safetensors>=0.4.4
piexif>=1.1.3

2
save.py
View File

@ -20,7 +20,7 @@ def main():
flux = Flux1(
model_config=ModelConfig.from_alias(args.model),
quantize_full_weights=args.quantize,
quantize=args.quantize,
)
flux.save_model(args.path)

14
src/flux_1/config/runtime_config.py
View File

@ -13,31 +13,31 @@ class RuntimeConfig:
self.sigmas = self._create_sigmas(config, model_config)
@property
def height(self):
def height(self) -> int:
return self.config.height
@property
def width(self):
def width(self) -> int:
return self.config.width
@property
def guidance(self):
def guidance(self) -> float:
return self.config.guidance
@property
def num_inference_steps(self):
def num_inference_steps(self) -> int:
return self.config.num_inference_steps
@property
def precision(self):
def precision(self) -> mx.Dtype:
return self.config.precision
@property
def num_train_steps(self):
def num_train_steps(self) -> int:
return self.model_config.num_train_steps
@staticmethod
def _create_sigmas(config, model):
def _create_sigmas(config, model) -> mx.array:
sigmas = RuntimeConfig._create_sigmas_values(config.num_inference_steps)
if model == ModelConfig.FLUX1_DEV:
sigmas = RuntimeConfig._shift_sigmas(sigmas, config.width, config.height)

37
src/flux_1/flux.py
View File

@ -1,8 +1,6 @@
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
@ -14,6 +12,7 @@ from flux_1.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from flux_1.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
from flux_1.models.transformer.transformer import Transformer
from flux_1.models.vae.vae import VAE
from flux_1.post_processing.image import Image
from flux_1.post_processing.image_util import ImageUtil
from flux_1.tokenizer.clip_tokenizer import TokenizerCLIP
from flux_1.tokenizer.t5_tokenizer import TokenizerT5
@ -26,13 +25,12 @@ class Flux1:
def __init__(
self,
model_config: ModelConfig,
quantize_full_weights: int | None = None,
quantize: int | None = None,
local_path: str | None = None,
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
):
self.model_config = model_config
self.quantize_full_weights = quantize_full_weights
# 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)
@ -53,20 +51,22 @@ class Flux1:
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)
self.bits = None
if quantize is not None or weights.quantization_level is not None:
self.bits = weights.quantization_level if weights.quantization_level is not None else quantize
nn.quantize(self.vae, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.bits)
nn.quantize(self.transformer, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 64, group_size=64, bits=self.bits)
nn.quantize(self.t5_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.bits)
nn.quantize(self.clip_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=self.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:
def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> Image:
# Create a new runtime config based on the model type and input parameters
config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.num_inference_steps))
# 1. Create the initial latents
latents = mx.random.normal(
@ -80,7 +80,7 @@ class Flux1:
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)):
for t in time_steps:
# 3.t Predict the noise
noise = self.transformer.predict(
t=t,
@ -97,10 +97,17 @@ class Flux1:
# Evaluate to enable progress tracking
mx.eval(latents)
# 5. Decode the latent array
# 5. Decode the latent array and return the image
latents = Flux1._unpack_latents(latents, config.height, config.width)
decoded = self.vae.decode(latents)
return ImageUtil.to_image(decoded)
return ImageUtil.to_image(
decoded_latents=decoded,
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict['elapsed'],
config=config,
)
@staticmethod
def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
@ -130,7 +137,7 @@ class Flux1:
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)})
mx.save_safetensors(str(path / f"{i}.safetensors"), weight, {"quantization_level": str(self.bits)})
def _split_weights(weights: dict, max_file_size_gb: int = 2) -> list:
# Copied from mlx-examples repo

120
src/flux_1/post_processing/image.py Normal file
View File

@ -0,0 +1,120 @@
import json
import logging
from pathlib import Path
import PIL.Image
import mlx.core as mx
import piexif
from flux_1.config.model_config import ModelConfig
log = logging.getLogger(__name__)
class Image:
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,
):
self.image = image
self.model_config = model_config
self.seed = seed
self.prompt = prompt
self.steps = steps
self.guidance = guidance
self.precision = precision
self.quantization = quantization
self.generation_time = generation_time
def save(self, path: str, export_json_metadata: bool = False) -> None:
file_path = Path(path)
file_path.parent.mkdir(parents=True, exist_ok=True)
file_name = file_path.stem
file_extension = file_path.suffix
# If a file already exists, create a new name with a counter
counter = 1
while file_path.exists():
new_name = f"{file_name}_{counter}{file_extension}"
file_path = file_path.with_name(new_name)
counter += 1
try:
# Save image without metadata first
self.image.save(file_path)
log.info(f"Image saved successfully at: {file_path}")
# Optionally save json metadata file
if export_json_metadata:
with open(f"{file_path.with_suffix(".json")}", 'w') as json_file:
json.dump(self._get_metadata(), json_file, indent=4)
# Embed metadata
self._embed_metadata(file_path)
log.info(f"Metadata embedded successfully at: {file_path}")
except Exception as e:
log.error(f"Error saving image: {e}")
def _embed_metadata(self, path: str) -> None:
try:
# Prepare metadata
metadata = self._get_metadata()
# Convert metadata dictionary to a string
metadata_str = str(metadata)
# Convert the string to bytes (using UTF-8 encoding)
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
}
}
# 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
# Save the image with metadata
image.save(path, exif=exif_bytes)
except Exception as e:
log.error(f"Error embedding metadata: {e}")
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",
'prompt': self.prompt,
}

55
src/flux_1/post_processing/image_util.py
View File

@ -1,22 +1,37 @@
import logging
from pathlib import Path
import PIL
from PIL import Image
import mlx.core as mx
import numpy as np
from PIL import Image
log = logging.getLogger(__name__)
from flux_1.config.runtime_config import RuntimeConfig
from flux_1.post_processing.image import Image
class ImageUtil:
@staticmethod
def to_image(decoded_latents: mx.array) -> PIL.Image.Image:
def to_image(
decoded_latents: mx.array,
seed: int,
prompt: str,
quantization: int,
generation_time: float,
config: RuntimeConfig,
) -> Image:
normalized = ImageUtil._denormalize(decoded_latents)
normalized_numpy = ImageUtil._to_numpy(normalized)
image = ImageUtil._numpy_to_pil(normalized_numpy)
return image
return Image(
image=image,
model_config=config.model_config,
seed=seed,
steps=config.num_inference_steps,
prompt=prompt,
guidance=config.guidance,
precision=config.precision,
quantization=quantization,
generation_time=generation_time,
)
@staticmethod
def _denormalize(images: mx.array) -> mx.array:
@ -36,7 +51,7 @@ class ImageUtil:
@staticmethod
def _numpy_to_pil(images: np.ndarray) -> PIL.Image.Image:
images = (images * 255).round().astype("uint8")
pil_images = [Image.fromarray(image) for image in images]
pil_images = [PIL.Image.fromarray(image) for image in images]
return pil_images[0]
@staticmethod
@ -47,7 +62,7 @@ class ImageUtil:
@staticmethod
def to_array(image: PIL.Image.Image) -> mx.array:
image = ImageUtil.resize(image)
image = ImageUtil._resize(image)
image = ImageUtil._pil_to_numpy(image)
array = mx.array(image)
array = mx.transpose(array, (0, 3, 1, 2))
@ -55,26 +70,6 @@ class ImageUtil:
return array
@staticmethod
def resize(image):
def _resize(image):
image = image.resize((1024, 1024), resample=PIL.Image.LANCZOS)
return image
@staticmethod
def save_image(image: Image.Image, path: str) -> None:
file_path = Path(path)
file_path.parent.mkdir(parents=True, exist_ok=True)
file_name = file_path.stem
file_extension = file_path.suffix
# If a file already exists, create a new name with a counter
counter = 1
while file_path.exists():
new_name = f"{file_name}({counter}){file_extension}"
file_path = file_path.with_name(new_name)
counter += 1
try:
image.save(file_path)
log.info(f"Image saved successfully at: {file_path}")
except Exception as e:
log.info(f"Error saving image: {e}")