112 lines
4.1 KiB
Python
112 lines
4.1 KiB
Python
from mflux.controlnet.transformer_controlnet import TransformerControlnet
|
|
from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet
|
|
from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
|
|
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
|
|
from mflux.models.transformer.transformer import Transformer
|
|
from mflux.models.vae.vae import VAE
|
|
from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
|
|
from mflux.tokenizer.t5_tokenizer import TokenizerT5
|
|
from mflux.tokenizer.tokenizer_handler import TokenizerHandler
|
|
from mflux.weights.weight_handler import WeightHandler
|
|
from mflux.weights.weight_handler_lora import WeightHandlerLoRA
|
|
from mflux.weights.weight_util import WeightUtil
|
|
|
|
|
|
class FluxInitializer:
|
|
@staticmethod
|
|
def init(
|
|
flux_model,
|
|
model_config,
|
|
quantize: int | None,
|
|
local_path: str | None,
|
|
lora_paths: list[str] | None,
|
|
lora_scales: list[float] | None,
|
|
) -> None:
|
|
# 0. Set paths and config for later
|
|
flux_model.lora_paths = lora_paths
|
|
flux_model.lora_scales = lora_scales
|
|
flux_model.model_config = model_config
|
|
|
|
# 1. Initialize tokenizers
|
|
tokenizers = TokenizerHandler(
|
|
repo_id=model_config.model_name,
|
|
max_t5_length=model_config.max_sequence_length,
|
|
local_path=local_path,
|
|
)
|
|
flux_model.t5_tokenizer = TokenizerT5(
|
|
tokenizer=tokenizers.t5,
|
|
max_length=model_config.max_sequence_length
|
|
) # fmt: off
|
|
flux_model.clip_tokenizer = TokenizerCLIP(
|
|
tokenizer=tokenizers.clip,
|
|
)
|
|
|
|
# 2. Load the regular weights
|
|
weights = WeightHandler.load_regular_weights(
|
|
repo_id=model_config.model_name,
|
|
local_path=local_path
|
|
) # fmt: off
|
|
|
|
# 3. Initialize all models
|
|
flux_model.vae = VAE()
|
|
flux_model.transformer = Transformer(
|
|
model_config=model_config,
|
|
num_transformer_blocks=weights.num_transformer_blocks(),
|
|
num_single_transformer_blocks=weights.num_single_transformer_blocks(),
|
|
)
|
|
flux_model.t5_text_encoder = T5Encoder()
|
|
flux_model.clip_text_encoder = CLIPEncoder()
|
|
|
|
# 4. Apply weights and quantize the models
|
|
flux_model.bits = WeightUtil.set_weights_and_quantize(
|
|
quantize_arg=quantize,
|
|
weights=weights,
|
|
vae=flux_model.vae,
|
|
transformer=flux_model.transformer,
|
|
t5_text_encoder=flux_model.t5_text_encoder,
|
|
clip_text_encoder=flux_model.clip_text_encoder,
|
|
)
|
|
|
|
# 5. Set LoRA weights
|
|
lora_weights = WeightHandlerLoRA.load_lora_weights(
|
|
transformer=flux_model.transformer,
|
|
lora_files=lora_paths,
|
|
lora_scales=lora_scales,
|
|
)
|
|
WeightHandlerLoRA.set_lora_weights(
|
|
transformer=flux_model.transformer,
|
|
loras=lora_weights
|
|
) # fmt: off
|
|
|
|
@staticmethod
|
|
def init_controlnet(
|
|
flux_model,
|
|
model_config,
|
|
quantize: int | None,
|
|
local_path: str | None,
|
|
lora_paths: list[str] | None,
|
|
lora_scales: list[float] | None,
|
|
) -> None:
|
|
# 1. Start with same init as regular Flux
|
|
FluxInitializer.init(
|
|
flux_model=flux_model,
|
|
model_config=model_config,
|
|
quantize=quantize,
|
|
local_path=local_path,
|
|
lora_paths=lora_paths,
|
|
lora_scales=lora_scales,
|
|
)
|
|
|
|
# 2. Apply ControlNet-specific initialization
|
|
weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer()
|
|
flux_model.transformer_controlnet = TransformerControlnet(
|
|
model_config=model_config,
|
|
num_transformer_blocks=weights_controlnet.num_transformer_blocks(),
|
|
num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks(),
|
|
)
|
|
WeightUtil.set_controlnet_weights_and_quantize(
|
|
quantize_arg=quantize,
|
|
weights=weights_controlnet,
|
|
transformer_controlnet=flux_model.transformer_controlnet,
|
|
)
|