Merge pull request #2 from filipstrand/automatically-download-model

Automatically download models from Huggingface or used cached ones

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. 

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")

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
+tqdm>=4.66.5
+huggingface-hub>=0.24.5

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)

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/**"
+                ]
+            )
+        )

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",
+                ]
+            )
+        )