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

Automatically download models from huggingface or used cache ones

Browse Source
This commit is contained in:
filipstrand 2024-08-15 00:17:00 +02:00
parent 97028713e1
commit e3a85df133
4 changed files with 58 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
requirements.txt
View File

@ -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

6
src/flux_1_schnell/models/flux.py
View File

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

27
src/flux_1_schnell/tokenizer/tokenizer_handler.py
View File

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

43
src/flux_1_schnell/weights/weight_handler.py
View File

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