From e3a85df1331ee931c7deccaf9e91e8f37c1067e8 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Thu, 15 Aug 2024 00:17:00 +0200 Subject: [PATCH 1/2] Automatically download models from huggingface or used cache ones --- requirements.txt | 3 +- src/flux_1_schnell/models/flux.py | 6 +-- .../tokenizer/tokenizer_handler.py | 27 +++++++++--- src/flux_1_schnell/weights/weight_handler.py | 43 +++++++++++++------ 4 files changed, 58 insertions(+), 21 deletions(-) diff --git a/requirements.txt b/requirements.txt index dbc01f7..2536076 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ pillow>=10.4.0 transformers>=4.44.0 sentencepiece>=0.2.0 torch>=2.3.1 -tqdm>=4.66.5 \ No newline at end of file +tqdm>=4.66.5 +huggingface-hub>=0.24.5 \ No newline at end of file diff --git a/src/flux_1_schnell/models/flux.py b/src/flux_1_schnell/models/flux.py index 3a7ab9d..4867027 100644 --- a/src/flux_1_schnell/models/flux.py +++ b/src/flux_1_schnell/models/flux.py @@ -21,12 +21,12 @@ from flux_1_schnell.weights.weight_handler import WeightHandler class Flux1Schnell: - def __init__(self, root_path: str): - tokenizers = TokenizerHandler.load_from_disk_via_huggingface_transformers(root_path) + def __init__(self, repo_id: str): + tokenizers = TokenizerHandler.load_from_disk_or_huggingface(repo_id) self.clip_tokenizer = TokenizerCLIP(tokenizers.clip) self.t5_tokenizer = TokenizerT5(tokenizers.t5) - weights = WeightHandler.load_from_disk(root_path) + weights = WeightHandler.load_from_disk_or_huggingface(repo_id) self.vae = VAE(weights.vae) self.transformer = Transformer(weights.transformer) self.clip_text_encoder = CLIPEncoder(weights.clip_encoder) diff --git a/src/flux_1_schnell/tokenizer/tokenizer_handler.py b/src/flux_1_schnell/tokenizer/tokenizer_handler.py index ff1e13a..8c754bb 100644 --- a/src/flux_1_schnell/tokenizer/tokenizer_handler.py +++ b/src/flux_1_schnell/tokenizer/tokenizer_handler.py @@ -1,4 +1,7 @@ +from pathlib import Path + import transformers +from huggingface_hub import snapshot_download from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5 @@ -6,18 +9,32 @@ from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5 class TokenizerHandler: - def __init__(self, root_path: str): + def __init__(self, repo_id: str): + root_path = TokenizerHandler._download_or_get_cached_tokenizers(repo_id) + self.clip = transformers.CLIPTokenizer.from_pretrained( - pretrained_model_name_or_path=root_path + "/tokenizer", + pretrained_model_name_or_path=root_path / "tokenizer", local_files_only=True, max_length=TokenizerCLIP.MAX_TOKEN_LENGTH ) self.t5 = transformers.T5Tokenizer.from_pretrained( - pretrained_model_name_or_path=root_path + "/tokenizer_2", + pretrained_model_name_or_path=root_path / "tokenizer_2", local_files_only=True, max_length=TokenizerT5.MAX_TOKEN_LENGTH ) @staticmethod - def load_from_disk_via_huggingface_transformers(root_path: str) -> "TokenizerHandler": - return TokenizerHandler(root_path) + def load_from_disk_or_huggingface(repo_id: str) -> "TokenizerHandler": + return TokenizerHandler(repo_id) + + @staticmethod + def _download_or_get_cached_tokenizers(repo_id: str) -> Path: + return Path( + snapshot_download( + repo_id=repo_id, + allow_patterns=[ + "tokenizer/**", + "tokenizer_2/**" + ] + ) + ) diff --git a/src/flux_1_schnell/weights/weight_handler.py b/src/flux_1_schnell/weights/weight_handler.py index c154e00..e9147c2 100644 --- a/src/flux_1_schnell/weights/weight_handler.py +++ b/src/flux_1_schnell/weights/weight_handler.py @@ -1,4 +1,7 @@ +from pathlib import Path + import mlx.core as mx +from huggingface_hub import snapshot_download from mlx.utils import tree_unflatten from flux_1_schnell.config.config import Config @@ -6,30 +9,32 @@ from flux_1_schnell.config.config import Config class WeightHandler: - def __init__(self, root_path: str): + def __init__(self, repo_id: str): + root_path = WeightHandler._download_or_get_cached_weights(repo_id) + self.clip_encoder = WeightHandler._clip_encoder( - weights=WeightHandler._load(root_path + "/text_encoder/model.safetensors") + 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"), + 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") + 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"), + 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"), ) @staticmethod - def load_from_disk(root_path: str) -> "WeightHandler": - return WeightHandler(root_path) + def load_from_disk_or_huggingface(repo_id: str) -> "WeightHandler": + return WeightHandler(repo_id) @staticmethod - def _load(path: str) -> list[dict]: - return list(mx.load(path).items()) + def _load(path: Path) -> list[dict]: + return list(mx.load(str(path)).items()) @staticmethod def _clip_encoder(weights: list[dict]) -> dict: @@ -108,3 +113,17 @@ class WeightHandler: value = value.transpose(0, 2, 3, 1) value = value.reshape(-1).reshape(value.shape).astype(Config.precision) return [(key, value)] + + @staticmethod + def _download_or_get_cached_weights(repo_id: str) -> Path: + return Path( + snapshot_download( + repo_id=repo_id, + allow_patterns=[ + "text_encoder/*.safetensors", + "text_encoder_2/*.safetensors", + "transformer/*.safetensors", + "vae/*.safetensors", + ] + ) + ) From f0b02257e3a38f1476c75ed61290b901ba17a4d5 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Thu, 15 Aug 2024 00:17:08 +0200 Subject: [PATCH 2/2] Update README with installation guide --- README.md | 40 +++++++++++++++++++-------- src/flux_1_schnell/main.py => main.py | 9 ++++-- 2 files changed, 36 insertions(+), 13 deletions(-) rename src/flux_1_schnell/main.py => main.py (78%) diff --git a/README.md b/README.md index e016347..7e3a408 100644 --- a/README.md +++ b/README.md @@ -22,23 +22,36 @@ like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) f - [x] FLUX.1-Scnhell - [ ] FLUX.1-Dev -### Prerequisites - -- The required libraries are listed in [requirements.txt](requirements.txt). -- Development and testing was done with Python 3.12 -- Download the weights and tokenizers at [Huggingface/black-forest-labs](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main). - Point to the location of the root folder as shown in the example below: - +### Installation +1. Clone the repo: + ``` + git clone git@github.com/filipstrand/mflux.git + ``` +2. Navigate to the project and set up a virtual environment: + ``` + cd mflux && python3 -m venv .venv && source .venv/bin/activate + ``` +3. Install the required dependencies: + ``` + pip install -r requirements.txt + ``` ### Generating an image -Run [main.py](/src/flux_1_schnell/main.py) or make a new script like the following -(make sure to correctly define the root path to wherever the model is saved and where you want the images to be stored). +Run the provided [main.py](main.py) +``` +python main.py +``` + +or make a new separate script like the following ```python +import sys +sys.path.append("/path/to/mflux/src") + from flux_1_schnell.config.config import Config from flux_1_schnell.models.flux import Flux1Schnell -flux = Flux1Schnell("/Users/filipstrand/.cache/FLUX.1-schnell/") +flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell") image = flux.generate_image( seed=3, @@ -48,13 +61,18 @@ image = flux.generate_image( ) ) -image.save(f"/Users/filipstrand/Desktop/image.png") +image.save("image.png") ``` +If the 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 model). + Generating a single image (with 2 inference steps, Schnell model) takes between 2 and 3 minutes. This implementation has been tested on two Macbook Pro machines: - 2021 M1 Pro (32 GB) - 2023 M2 Max (32 GB) +Update: +On faster machines, [@karpathy](https://gist.github.com/awni/a67d16d50f0f492d94a10418e0592bde?permalink_comment_id=5153531#gistcomment-5153531) and [@awni](https://x.com/awnihannun/status/1823515121827897385) have reported times ~20s and below! + ### Equivalent to Diffusers implementation There is only a single source of randomness when generating an image: The initial latent array. diff --git a/src/flux_1_schnell/main.py b/main.py similarity index 78% rename from src/flux_1_schnell/main.py rename to main.py index 7015544..bc3c036 100644 --- a/src/flux_1_schnell/main.py +++ b/main.py @@ -1,7 +1,12 @@ +import os +import sys + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src'))) + from flux_1_schnell.config.config import Config from flux_1_schnell.models.flux import Flux1Schnell -flux = Flux1Schnell("/Users/filipstrand/.cache/FLUX.1-schnell/") +flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell") image = flux.generate_image( seed=3, @@ -11,4 +16,4 @@ image = flux.generate_image( ) ) -image.save("/Users/filipstrand/Desktop/image.png") +image.save("image.png")