Spaces:
Paused
Paused
File size: 970 Bytes
99b955f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
from pathlib import Path
from constants import DEVICE, LCM_DEFAULT_MODEL_OPENVINO, TAEF1_MODEL_OPENVINO
from huggingface_hub import snapshot_download
from backend.openvino.ovflux import (
TEXT_ENCODER_2_PATH,
TEXT_ENCODER_PATH,
TRANSFORMER_PATH,
VAE_DECODER_PATH,
init_pipeline,
)
def get_flux_pipeline(
model_id: str = LCM_DEFAULT_MODEL_OPENVINO,
use_taef1: bool = False,
taef1_path: str = TAEF1_MODEL_OPENVINO,
):
model_dir = Path(snapshot_download(model_id))
vae_dir = Path(snapshot_download(taef1_path)) if use_taef1 else model_dir
model_dict = {
"transformer": model_dir / TRANSFORMER_PATH,
"text_encoder": model_dir / TEXT_ENCODER_PATH,
"text_encoder_2": model_dir / TEXT_ENCODER_2_PATH,
"vae": vae_dir / VAE_DECODER_PATH,
}
ov_pipe = init_pipeline(
model_dir,
model_dict,
device=DEVICE.upper(),
use_taef1=use_taef1,
)
return ov_pipe
|