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