diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index d26568cb5..edd634e79 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -7,7 +7,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -32,8 +31,6 @@ def _preparation_modules(self, config: TrainConfig, model: ErnieModel): image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean') downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125) tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673, unclear if Mistral is affected encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=HIDDEN_STATES_LAYER, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index 587f480c5..0aa4522d7 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -14,7 +14,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -43,8 +42,6 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length, apply_chat_template = lambda caption: mistral_format_input([caption], MISTRAL_SYSTEM_MESSAGE), apply_chat_template_kwargs = {'add_generation_prompt': False}, ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(), hidden_state_output_index=MISTRAL_HIDDEN_STATES_LAYERS, @@ -53,8 +50,6 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length, apply_chat_template = lambda caption: qwen3_format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': False} ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=QWEN3_HIDDEN_STATES_LAYERS, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py index 03934dd4d..66cfea16f 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -9,7 +9,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -38,8 +37,6 @@ def _preparation_modules(self, config: TrainConfig, model: ZImageModel): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH, apply_chat_template = lambda caption: format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': True} ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=-2, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) prune_masked_tokens = PruneMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state') diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index 19146bd2b..59967c8dc 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -40,6 +40,7 @@ def __init__( class ChromaModel(BaseModel): # base model data tokenizer: T5Tokenizer | None + orig_tokenizer: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder: T5EncoderModel | None vae: AutoencoderKL | None @@ -72,6 +73,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -141,13 +143,13 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return ChromaPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=self.orig_tokenizer if use_original_tokenizers else self.tokenizer, ) def add_text_encoder_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index e4bfaa066..b981865c4 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -50,7 +50,9 @@ def __init__( class FluxModel(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: T5Tokenizer | None + orig_tokenizer_2: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder_1: CLIPTextModel | None text_encoder_2: T5EncoderModel | None @@ -86,7 +88,9 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -177,15 +181,15 @@ def eval(self): self.text_encoder_2.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return FluxPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index 121f5fd7c..049b0e6d3 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -267,19 +267,19 @@ def eval(self): self.text_encoder_4.eval() self.transformer.eval() - def create_pipeline(self, use_original_modules: bool) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return HiDreamImagePipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.orig_tokenizer_1 if use_original_modules else self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.orig_tokenizer_2 if use_original_modules else self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, text_encoder_3=self.text_encoder_3, - tokenizer_3=self.orig_tokenizer_3 if use_original_modules else self.tokenizer_3, + tokenizer_3=self.orig_tokenizer_3 if use_original_tokenizers else self.tokenizer_3, text_encoder_4=self.text_encoder_4, - tokenizer_4=self.orig_tokenizer_4 if use_original_modules else self.tokenizer_4, + tokenizer_4=self.orig_tokenizer_4 if use_original_tokenizers else self.tokenizer_4, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index 49c40cc15..2107e55d7 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -196,15 +196,15 @@ def eval(self): self.text_encoder_2.eval() self.transformer.eval() - def create_pipeline(self, use_original_modules: bool) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return HunyuanVideoPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.orig_tokenizer_1 if use_original_modules else self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.orig_tokenizer_2 if use_original_modules else self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index 97f87e5ca..466cc61f9 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -42,6 +42,7 @@ def __init__( class PixArtAlphaModel(BaseModel): # base model data tokenizer: T5Tokenizer | None + orig_tokenizer: T5Tokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: T5EncoderModel | None vae: AutoencoderKL | None @@ -74,6 +75,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -141,11 +143,12 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + tokenizer = self.orig_tokenizer if use_original_tokenizers else self.tokenizer match self.model_type: case ModelType.PIXART_ALPHA: return PixArtAlphaPipeline( - tokenizer=self.tokenizer, + tokenizer=tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, @@ -153,7 +156,7 @@ def create_pipeline(self) -> DiffusionPipeline: ) case ModelType.PIXART_SIGMA: return PixArtSigmaPipeline( - tokenizer=self.tokenizer, + tokenizer=tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index 506afe741..9e8008219 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -41,6 +41,7 @@ def __init__( class SanaModel(BaseModel): # base model data tokenizer: GemmaTokenizer | None + orig_tokenizer: GemmaTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: Gemma2Model | None vae: AutoencoderDC | None @@ -74,6 +75,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -143,9 +145,9 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return SanaPipeline( - tokenizer=self.tokenizer, + tokenizer=self.orig_tokenizer if use_original_tokenizers else self.tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index bcaf46fcb..8f6cf5818 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -58,8 +58,11 @@ def __init__( class StableDiffusion3Model(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: CLIPTokenizer | None + orig_tokenizer_2: CLIPTokenizer | None tokenizer_3: T5Tokenizer | None + orig_tokenizer_3: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder_1: CLIPTextModelWithProjection | None text_encoder_2: CLIPTextModelWithProjection | None @@ -98,8 +101,11 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.tokenizer_3 = None + self.orig_tokenizer_3 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -208,17 +214,17 @@ def eval(self): self.text_encoder_3.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return StableDiffusion3Pipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, text_encoder_3=self.text_encoder_3, - tokenizer_3=self.tokenizer_3, + tokenizer_3=self.orig_tokenizer_3 if use_original_tokenizers else self.tokenizer_3, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/StableDiffusionModel.py b/modules/model/StableDiffusionModel.py index 2bee0fd70..4b1b11a2c 100644 --- a/modules/model/StableDiffusionModel.py +++ b/modules/model/StableDiffusionModel.py @@ -43,6 +43,7 @@ def __init__( class StableDiffusionModel(BaseModel): # base model data tokenizer: CLIPTokenizer | None + orig_tokenizer: CLIPTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: CLIPTextModel | None vae: AutoencoderKL | None @@ -72,6 +73,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -136,12 +138,13 @@ def eval(self): self.text_encoder.eval() self.unet.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + tokenizer = self.orig_tokenizer if use_original_tokenizers else self.tokenizer if self.model_type.has_depth_input(): return StableDiffusionDepth2ImgPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, depth_estimator=self.depth_estimator, @@ -151,7 +154,7 @@ def create_pipeline(self) -> DiffusionPipeline: return StableDiffusionInpaintPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, safety_checker=None, @@ -162,7 +165,7 @@ def create_pipeline(self) -> DiffusionPipeline: return StableDiffusionPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, safety_checker=None, diff --git a/modules/model/StableDiffusionXLModel.py b/modules/model/StableDiffusionXLModel.py index b68df2d23..43edead8e 100644 --- a/modules/model/StableDiffusionXLModel.py +++ b/modules/model/StableDiffusionXLModel.py @@ -45,7 +45,9 @@ def __init__( class StableDiffusionXLModel(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: CLIPTokenizer | None + orig_tokenizer_2: CLIPTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder_1: CLIPTextModel | None text_encoder_2: CLIPTextModelWithProjection | None @@ -81,7 +83,9 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -166,13 +170,13 @@ def eval(self): self.text_encoder_2.eval() self.unet.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return StableDiffusionXLPipeline( vae=self.vae, text_encoder=self.text_encoder_1, text_encoder_2=self.text_encoder_2, - tokenizer=self.tokenizer_1, - tokenizer_2=self.tokenizer_2, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, unet=self.unet, scheduler=self.noise_scheduler, ) diff --git a/modules/model/WuerstchenModel.py b/modules/model/WuerstchenModel.py index b844d6eb1..7b50d81ed 100644 --- a/modules/model/WuerstchenModel.py +++ b/modules/model/WuerstchenModel.py @@ -69,6 +69,7 @@ class WuerstchenModel(BaseModel): decoder_vqgan: PaellaVQModel | None effnet_encoder: WuerstchenEfficientNetEncoder | None prior_tokenizer: CLIPTokenizer | None + orig_prior_tokenizer: CLIPTokenizer | None prior_text_encoder: CLIPTextModel | None prior_noise_scheduler: DDPMWuerstchenScheduler | None prior_prior: WuerstchenPrior | StableCascadeUNet | None @@ -105,6 +106,7 @@ def __init__( self.decoder_vqgan = None self.effnet_encoder = None self.prior_tokenizer = None + self.orig_prior_tokenizer = None self.prior_text_encoder = None self.prior_noise_scheduler = None self.prior_prior = None @@ -179,7 +181,8 @@ def eval(self): self.prior_text_encoder.eval() self.prior_prior.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + prior_tokenizer = self.orig_prior_tokenizer if use_original_tokenizers else self.prior_tokenizer if self.model_type.is_wuerstchen_v2(): return WuerstchenCombinedPipeline( tokenizer=self.decoder_tokenizer, @@ -187,19 +190,19 @@ def create_pipeline(self) -> DiffusionPipeline: decoder=self.decoder_decoder, scheduler=self.decoder_noise_scheduler, vqgan=self.decoder_vqgan, - prior_tokenizer=self.prior_tokenizer, + prior_tokenizer=prior_tokenizer, prior_text_encoder=self.prior_text_encoder, prior_prior=self.prior_prior, prior_scheduler=self.prior_noise_scheduler, ) elif self.model_type.is_stable_cascade(): return StableCascadeCombinedPipeline( - tokenizer=self.prior_tokenizer, + tokenizer=prior_tokenizer, text_encoder=self.prior_text_encoder, decoder=self.decoder_decoder, scheduler=self.decoder_noise_scheduler, vqgan=self.decoder_vqgan, - prior_tokenizer=self.prior_tokenizer, + prior_tokenizer=prior_tokenizer, prior_text_encoder=self.prior_text_encoder, prior_prior=self.prior_prior, prior_scheduler=self.prior_noise_scheduler, diff --git a/modules/modelLoader/ErnieModelLoader.py b/modules/modelLoader/ErnieModelLoader.py index 8e8365981..c9d623463 100644 --- a/modules/modelLoader/ErnieModelLoader.py +++ b/modules/modelLoader/ErnieModelLoader.py @@ -1,4 +1,3 @@ -import logging import os import traceback @@ -22,8 +21,7 @@ FlowMatchEulerDiscreteScheduler, GGUFQuantizationConfig, ) -from transformers import Mistral3Model, MistralConfig, PreTrainedTokenizerFast -from transformers.models.auto.configuration_auto import CONFIG_MAPPING +from transformers import AutoTokenizer, Mistral3Model class ErnieModelLoader( @@ -60,23 +58,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - # transformers < 5.x doesn't register "ministral3"; patch it so Mistral3Config can parse its text_config - if "ministral3" not in CONFIG_MAPPING: - CONFIG_MAPPING.register("ministral3", MistralConfig) - - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if transformer_model_name: transformer = ErnieImageTransformer2DModel.from_single_file( transformer_model_name, @@ -98,15 +79,10 @@ def __load_diffusers( quantization, ) - # TokenizersBackend is the Rust tokenizers library backend, not a transformers class — warning is a false alarm - tokenization_logger = logging.getLogger("transformers.tokenization_utils_base") - prev_level = tokenization_logger.level - tokenization_logger.setLevel(logging.ERROR) - tokenizer = PreTrainedTokenizerFast.from_pretrained( + tokenizer = AutoTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - tokenization_logger.setLevel(prev_level) text_encoder = self._load_transformers_sub_module( Mistral3Model, diff --git a/modules/modelLoader/Flux2ModelLoader.py b/modules/modelLoader/Flux2ModelLoader.py index 414e43232..55bbeebd4 100644 --- a/modules/modelLoader/Flux2ModelLoader.py +++ b/modules/modelLoader/Flux2ModelLoader.py @@ -62,19 +62,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if transformer_model_name: transformer = Flux2Transformer2DModel.from_single_file( transformer_model_name, @@ -122,10 +109,6 @@ def __load_diffusers( base_model_name, "text_encoder", ) - #TODO this is a tied weight. The dtype conversion code in _load_transformers_sub_module - #currently does not support tied weights. Reconstruct but clone, because the quantization code - #doesn't support tied weights either: - text_encoder.lm_head.weight = type(text_encoder.lm_head.weight)(text_encoder.model.embed_tokens.weight) noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, diff --git a/modules/modelLoader/ZImageModelLoader.py b/modules/modelLoader/ZImageModelLoader.py index 0fe6e3c00..aee70fbe4 100644 --- a/modules/modelLoader/ZImageModelLoader.py +++ b/modules/modelLoader/ZImageModelLoader.py @@ -60,19 +60,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", @@ -91,12 +78,6 @@ def __load_diffusers( "text_encoder", ) - #TODO this is a tied weight. The dtype conversion code in _load_transformers_sub_module - #currently does not support tied weights. Reconstruct but clone, because the quantization code - #doesn't support tied weights either: - text_encoder.lm_head.weight = type(text_encoder.lm_head.weight)(text_encoder.model.embed_tokens.weight) - - if vae_model_name: vae = self._load_diffusers_sub_module( AutoencoderKL, diff --git a/modules/modelLoader/chroma/ChromaModelLoader.py b/modules/modelLoader/chroma/ChromaModelLoader.py index 9c5a5d543..7dcbef794 100644 --- a/modules/modelLoader/chroma/ChromaModelLoader.py +++ b/modules/modelLoader/chroma/ChromaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -52,18 +53,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=["text_encoder"], - ) - tokenizer = T5Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", @@ -120,6 +109,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/flux/FluxModelLoader.py b/modules/modelLoader/flux/FluxModelLoader.py index 7747a9f3e..d4f21ea2a 100644 --- a/modules/modelLoader/flux/FluxModelLoader.py +++ b/modules/modelLoader/flux/FluxModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -58,23 +59,6 @@ def __load_diffusers( include_text_encoder_2: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if include_text_encoder_1: tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, @@ -156,7 +140,9 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -231,7 +217,9 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = pipeline.scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/hiDream/HiDreamModelLoader.py b/modules/modelLoader/hiDream/HiDreamModelLoader.py index 639156a38..b3e20f23c 100644 --- a/modules/modelLoader/hiDream/HiDreamModelLoader.py +++ b/modules/modelLoader/hiDream/HiDreamModelLoader.py @@ -67,34 +67,6 @@ def __load_diffusers( include_text_encoder_4: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if include_text_encoder_3: - transformers_sub.append("text_encoder_3") - if include_text_encoder_4: - if text_encoder_4_model_name: - self._prepare_sub_modules( - text_encoder_4_model_name, - transformers_modules=[""], - diffusers_modules=[], - ) - else: - transformers_sub.append("text_encoder_4") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", diff --git a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py index 0e9d577cc..85c91699b 100644 --- a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py +++ b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py @@ -59,24 +59,6 @@ def __load_diffusers( include_text_encoder_2: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - - if not transformer_model_name: - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if include_text_encoder_1: tokenizer_1 = LlamaTokenizerFast.from_pretrained( base_model_name, diff --git a/modules/modelLoader/mixin/HFModelLoaderMixin.py b/modules/modelLoader/mixin/HFModelLoaderMixin.py index 9180014bb..f2f196257 100644 --- a/modules/modelLoader/mixin/HFModelLoaderMixin.py +++ b/modules/modelLoader/mixin/HFModelLoaderMixin.py @@ -1,7 +1,6 @@ import json +import logging import os -import re -import traceback from abc import ABCMeta from itertools import repeat @@ -15,11 +14,17 @@ import torch from torch import nn +from transformers.conversion_mapping import get_checkpoint_conversion_mapping +from transformers.core_model_loading import rename_source_key + import accelerate import huggingface_hub from huggingface_hub.utils import EntryNotFoundError from safetensors.torch import load_file +# huggingface_hub 1.16+ uses httpx, which logs every HTTP request/response at INFO level. +logging.getLogger("httpx").setLevel(logging.WARNING) + class HFModelLoaderMixin(metaclass=ABCMeta): def __init__(self): @@ -114,12 +119,20 @@ def __load_sub_module( if hasattr(sub_module, '_fix_state_dict_keys_on_load'): sub_module._fix_state_dict_keys_on_load(state_dict) - if hasattr(sub_module, "_checkpoint_conversion_mapping"): #required for loading the text encoder of Qwen + #some checkpoints (e.g. Ernie's Mistral3 text encoder, Qwen's Qwen2_5_VL text encoder) were saved with an + #older module layout than the one transformers builds from the config in this version. transformers' own + #from_pretrained applies the same renaming via its checkpoint conversion registry, so we reuse it here. + #diffusers sub-modules have no such registry (their config is a plain FrozenDict, no model_type), and + #never need this renaming. + weight_renamings = get_checkpoint_conversion_mapping(sub_module.config.model_type) \ + if hasattr(sub_module.config, 'model_type') else None + if weight_renamings: + meta_state_dict = sub_module.state_dict() new_state_dict = {} for k, v in state_dict.items(): - new_k = k - for pattern, replacement in sub_module._checkpoint_conversion_mapping.items(): - new_k = re.sub(pattern, replacement, new_k) + new_k, _ = rename_source_key( + k, weight_renamings, [], prefix=sub_module.base_model_prefix, meta_state_dict=meta_state_dict, + ) new_state_dict[new_k] = v state_dict = new_state_dict @@ -158,6 +171,23 @@ def __load_sub_module( del state_dict + #tied weights (e.g. T5EncoderModel's encoder.embed_tokens.weight <-> shared.weight, or + #Qwen3ForCausalLM's lm_head.weight <-> model.embed_tokens.weight) are saved only once in the checkpoint, + #so the tied key above is never assigned and stays an empty meta tensor. populate it by cloning the + #source weight that was actually loaded, rather than aliasing the same Parameter object: aliasing would + #make a later in-place quantization of one side (e.g. quantize_layers() quantizing lm_head) silently + #corrupt the other (e.g. the embedding table), since both attribute paths would refer to the same object. + tied_weights_keys = getattr(sub_module, '_tied_weights_keys', None) + if tied_weights_keys is not None: + for target_key, source_key in tied_weights_keys.items(): + module = sub_module + *parents, tensor_name = target_key.split(".") + for p in parents: + module = getattr(module, p) + if module._parameters[tensor_name].is_meta: + source = sub_module.get_parameter(source_key) + module._parameters[tensor_name] = type(module._parameters[tensor_name])(source) + return sub_module def _load_transformers_sub_module( @@ -298,22 +328,3 @@ def _convert_diffusers_sub_module_to_dtype( None, quantization, ) - - def _prepare_sub_modules(self, pretrained_model_name_or_path: str, diffusers_modules: list[str], transformers_modules: list[str]): - is_local = os.path.isdir(pretrained_model_name_or_path) - if is_local: - return - - diffusers_paths = [((folder + "/") if folder else "") + "diffusion_pytorch_model*" for folder in diffusers_modules] - transformers_paths = [((folder + "/") if folder else "") + "model*" for folder in transformers_modules] - transformers_paths.extend([((folder + "/") if folder else "") + "pytorch_model*" for folder in transformers_modules]) - try: - huggingface_hub.snapshot_download( - pretrained_model_name_or_path, - allow_patterns=diffusers_paths + transformers_paths, - ) - except huggingface_hub.errors.HFValidationError: - pass - except Exception: - traceback.print_exc() - print("Error during bulk preloading of Huggingface model repository, proceeding without preloading") diff --git a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py index f8e75907d..467c29c8a 100644 --- a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py +++ b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -86,6 +87,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/qwen/QwenModelLoader.py b/modules/modelLoader/qwen/QwenModelLoader.py index 498b925af..953f15bfb 100644 --- a/modules/modelLoader/qwen/QwenModelLoader.py +++ b/modules/modelLoader/qwen/QwenModelLoader.py @@ -52,18 +52,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=["text_encoder"], - ) - tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", diff --git a/modules/modelLoader/sana/SanaModelLoader.py b/modules/modelLoader/sana/SanaModelLoader.py index 508579ba2..a904e3996 100644 --- a/modules/modelLoader/sana/SanaModelLoader.py +++ b/modules/modelLoader/sana/SanaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -86,6 +87,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py index 4333897ae..aa610f485 100644 --- a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py +++ b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -132,6 +133,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae @@ -206,6 +208,7 @@ def __load_ckpt( model.model_type = model_type model.tokenizer = pipeline.tokenizer + model.orig_tokenizer = copy.deepcopy(pipeline.tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae @@ -262,6 +265,7 @@ def __load_safetensors( model.model_type = model_type model.tokenizer = pipeline.tokenizer + model.orig_tokenizer = copy.deepcopy(pipeline.tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py index 17e55696d..47d87da74 100644 --- a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py +++ b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -50,8 +51,6 @@ def __load_diffusers( include_text_encoder_3: bool, quantization: QuantizationConfig, ): - #no call to self._prepare_sub_modules, because SAI polluted their sd3 / sd3.5 medium repo text encoders with fp16 files - if include_text_encoder_1: tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, @@ -141,8 +140,11 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.tokenizer_3 = tokenizer_3 + model.orig_tokenizer_3 = copy.deepcopy(tokenizer_3) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -230,8 +232,11 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.tokenizer_3 = tokenizer_3 + model.orig_tokenizer_3 = copy.deepcopy(tokenizer_3) model.noise_scheduler = pipeline.scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py index 6d31516fe..afbab6581 100644 --- a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py +++ b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -125,7 +126,9 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -167,7 +170,9 @@ def __load_ckpt( model.model_type = model_type model.tokenizer_1 = pipeline.tokenizer + model.orig_tokenizer_1 = copy.deepcopy(pipeline.tokenizer) model.tokenizer_2 = pipeline.tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(pipeline.tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -227,7 +232,9 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = pipeline.tokenizer + model.orig_tokenizer_1 = copy.deepcopy(pipeline.tokenizer) model.tokenizer_2 = pipeline.tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(pipeline.tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py index 8fec87b29..188107a2c 100644 --- a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py +++ b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py @@ -1,3 +1,4 @@ +import copy import json import os.path import traceback @@ -192,6 +193,7 @@ def __load_diffusers( model.decoder_vqgan = decoder_vqgan model.effnet_encoder = effnet_encoder model.prior_tokenizer = prior_tokenizer + model.orig_prior_tokenizer = copy.deepcopy(prior_tokenizer) model.prior_text_encoder = prior_text_encoder model.prior_noise_scheduler = prior_noise_scheduler model.prior_prior = prior_prior diff --git a/modules/modelSampler/HiDreamSampler.py b/modules/modelSampler/HiDreamSampler.py index d02197a35..dbafdae81 100644 --- a/modules/modelSampler/HiDreamSampler.py +++ b/modules/modelSampler/HiDreamSampler.py @@ -31,7 +31,7 @@ def __init__( self.model = model self.model_type = model_type - self.pipeline = model.create_pipeline(use_original_modules=False) + self.pipeline = model.create_pipeline() @torch.no_grad() def __sample_base( diff --git a/modules/modelSampler/HunyuanVideoSampler.py b/modules/modelSampler/HunyuanVideoSampler.py index 275626341..020df6a39 100644 --- a/modules/modelSampler/HunyuanVideoSampler.py +++ b/modules/modelSampler/HunyuanVideoSampler.py @@ -32,7 +32,7 @@ def __init__( self.model = model self.model_type = model_type - self.pipeline = model.create_pipeline(use_original_modules=False) + self.pipeline = model.create_pipeline() @torch.no_grad() def __sample_base( diff --git a/modules/modelSaver/chroma/ChromaModelSaver.py b/modules/modelSaver/chroma/ChromaModelSaver.py index f8129558d..2dee69bfc 100644 --- a/modules/modelSaver/chroma/ChromaModelSaver.py +++ b/modules/modelSaver/chroma/ChromaModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/flux/FluxModelSaver.py b/modules/modelSaver/flux/FluxModelSaver.py index dc632f2ae..b5ac966d5 100644 --- a/modules/modelSaver/flux/FluxModelSaver.py +++ b/modules/modelSaver/flux/FluxModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/hidream/HiDreamModelSaver.py b/modules/modelSaver/hidream/HiDreamModelSaver.py index 4bfe56ebc..b243166ec 100644 --- a/modules/modelSaver/hidream/HiDreamModelSaver.py +++ b/modules/modelSaver/hidream/HiDreamModelSaver.py @@ -24,7 +24,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline(use_original_modules=True) + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py b/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py index 866d68617..3fad1c7a0 100644 --- a/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py +++ b/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline(use_original_modules=True) + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py b/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py index 83e8af4a9..d6ef4d580 100644 --- a/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py +++ b/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py @@ -25,7 +25,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/sana/SanaModelSaver.py b/modules/modelSaver/sana/SanaModelSaver.py index 1dede8c8f..29301adbb 100644 --- a/modules/modelSaver/sana/SanaModelSaver.py +++ b/modules/modelSaver/sana/SanaModelSaver.py @@ -22,7 +22,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py b/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py index e0131f52f..f24e6b5f9 100644 --- a/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py +++ b/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py b/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py index bcfabab67..095b9e75f 100644 --- a/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py +++ b/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py b/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py index 09170e67b..24625fc52 100644 --- a/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py +++ b/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py @@ -26,7 +26,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py b/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py index 3bbe1b55d..cee5e28b4 100644 --- a/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py +++ b/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py @@ -25,7 +25,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline().prior_pipe + pipeline = model.create_pipeline(use_original_tokenizers=True).prior_pipe original_device = pipeline.device pipeline.to("cpu") pipeline_copy = copy.deepcopy(pipeline) diff --git a/modules/modelSetup/ChromaEmbeddingSetup.py b/modules/modelSetup/ChromaEmbeddingSetup.py index 9fa65b69f..63a6d75af 100644 --- a/modules/modelSetup/ChromaEmbeddingSetup.py +++ b/modules/modelSetup/ChromaEmbeddingSetup.py @@ -61,7 +61,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/ChromaFineTuneSetup.py b/modules/modelSetup/ChromaFineTuneSetup.py index 371c414ec..c78e4e700 100644 --- a/modules/modelSetup/ChromaFineTuneSetup.py +++ b/modules/modelSetup/ChromaFineTuneSetup.py @@ -70,7 +70,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/ChromaLoRASetup.py b/modules/modelSetup/ChromaLoRASetup.py index 1396b4a52..514eca8e8 100644 --- a/modules/modelSetup/ChromaLoRASetup.py +++ b/modules/modelSetup/ChromaLoRASetup.py @@ -98,7 +98,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxEmbeddingSetup.py b/modules/modelSetup/FluxEmbeddingSetup.py index 0418390da..98dd12579 100644 --- a/modules/modelSetup/FluxEmbeddingSetup.py +++ b/modules/modelSetup/FluxEmbeddingSetup.py @@ -71,8 +71,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxFineTuneSetup.py b/modules/modelSetup/FluxFineTuneSetup.py index a2dc78057..3afbdf31c 100644 --- a/modules/modelSetup/FluxFineTuneSetup.py +++ b/modules/modelSetup/FluxFineTuneSetup.py @@ -80,8 +80,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxLoRASetup.py b/modules/modelSetup/FluxLoRASetup.py index e3de941fd..e6204565f 100644 --- a/modules/modelSetup/FluxLoRASetup.py +++ b/modules/modelSetup/FluxLoRASetup.py @@ -122,8 +122,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py index 5d94c076d..5b842e267 100644 --- a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py +++ b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py @@ -58,7 +58,6 @@ def setup_model( ): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaFineTuneSetup.py b/modules/modelSetup/PixArtAlphaFineTuneSetup.py index 277f39004..31b92378e 100644 --- a/modules/modelSetup/PixArtAlphaFineTuneSetup.py +++ b/modules/modelSetup/PixArtAlphaFineTuneSetup.py @@ -74,7 +74,6 @@ def setup_model( if config.train_any_embedding(): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaLoRASetup.py b/modules/modelSetup/PixArtAlphaLoRASetup.py index 5c865b224..1e18dc673 100644 --- a/modules/modelSetup/PixArtAlphaLoRASetup.py +++ b/modules/modelSetup/PixArtAlphaLoRASetup.py @@ -94,7 +94,6 @@ def setup_model( model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.transformer_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaEmbeddingSetup.py b/modules/modelSetup/SanaEmbeddingSetup.py index 7704c3afd..c31c2aa11 100644 --- a/modules/modelSetup/SanaEmbeddingSetup.py +++ b/modules/modelSetup/SanaEmbeddingSetup.py @@ -58,7 +58,6 @@ def setup_model( ): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaFineTuneSetup.py b/modules/modelSetup/SanaFineTuneSetup.py index 11faa747a..1c334c44f 100644 --- a/modules/modelSetup/SanaFineTuneSetup.py +++ b/modules/modelSetup/SanaFineTuneSetup.py @@ -68,7 +68,6 @@ def setup_model( if config.train_any_embedding(): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaLoRASetup.py b/modules/modelSetup/SanaLoRASetup.py index f481cd445..45cb2f7bb 100644 --- a/modules/modelSetup/SanaLoRASetup.py +++ b/modules/modelSetup/SanaLoRASetup.py @@ -94,7 +94,6 @@ def setup_model( model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.transformer_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py index 1ff229928..e9209bdc4 100644 --- a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py @@ -81,9 +81,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3FineTuneSetup.py b/modules/modelSetup/StableDiffusion3FineTuneSetup.py index 3bbe41821..3ce400a59 100644 --- a/modules/modelSetup/StableDiffusion3FineTuneSetup.py +++ b/modules/modelSetup/StableDiffusion3FineTuneSetup.py @@ -90,9 +90,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3LoRASetup.py b/modules/modelSetup/StableDiffusion3LoRASetup.py index ae326cbbe..0d802c18a 100644 --- a/modules/modelSetup/StableDiffusion3LoRASetup.py +++ b/modules/modelSetup/StableDiffusion3LoRASetup.py @@ -148,9 +148,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionEmbeddingSetup.py b/modules/modelSetup/StableDiffusionEmbeddingSetup.py index 6e3a41c2b..f8cc956f2 100644 --- a/modules/modelSetup/StableDiffusionEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionEmbeddingSetup.py @@ -62,7 +62,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionFineTuneSetup.py b/modules/modelSetup/StableDiffusionFineTuneSetup.py index 8c42e59df..9f6fff8f3 100644 --- a/modules/modelSetup/StableDiffusionFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionFineTuneSetup.py @@ -76,7 +76,6 @@ def setup_model( elif config.force_epsilon_prediction: model.force_epsilon_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionLoRASetup.py b/modules/modelSetup/StableDiffusionLoRASetup.py index 9b9e9a613..d82147611 100644 --- a/modules/modelSetup/StableDiffusionLoRASetup.py +++ b/modules/modelSetup/StableDiffusionLoRASetup.py @@ -99,7 +99,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py index 7d9ef442f..62f78cba2 100644 --- a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py @@ -71,8 +71,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py index ec6dbe16c..56a5bfe40 100644 --- a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py @@ -86,8 +86,6 @@ def setup_model( elif config.force_epsilon_prediction: model.force_epsilon_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLLoRASetup.py b/modules/modelSetup/StableDiffusionXLLoRASetup.py index 6b193edc5..a0db29b07 100644 --- a/modules/modelSetup/StableDiffusionXLLoRASetup.py +++ b/modules/modelSetup/StableDiffusionXLLoRASetup.py @@ -120,8 +120,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenEmbeddingSetup.py b/modules/modelSetup/WuerstchenEmbeddingSetup.py index 37dc7f85f..d5dd8a83b 100644 --- a/modules/modelSetup/WuerstchenEmbeddingSetup.py +++ b/modules/modelSetup/WuerstchenEmbeddingSetup.py @@ -62,7 +62,6 @@ def setup_model( ): model.prior_text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenFineTuneSetup.py b/modules/modelSetup/WuerstchenFineTuneSetup.py index 1291125dc..111ce016f 100644 --- a/modules/modelSetup/WuerstchenFineTuneSetup.py +++ b/modules/modelSetup/WuerstchenFineTuneSetup.py @@ -71,7 +71,6 @@ def setup_model( if config.train_any_embedding(): model.prior_text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenLoRASetup.py b/modules/modelSetup/WuerstchenLoRASetup.py index b26f57436..812bb43e6 100644 --- a/modules/modelSetup/WuerstchenLoRASetup.py +++ b/modules/modelSetup/WuerstchenLoRASetup.py @@ -98,7 +98,6 @@ def setup_model( model.prior_prior_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.prior_prior_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py b/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py index 33b6d34cd..24f726230 100644 --- a/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py +++ b/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py @@ -13,27 +13,15 @@ CLIPTextModelWithProjection, Gemma2Model, LlamaModel, + PreTrainedTokenizer, T5EncoderModel, ) -from transformers.tokenization_utils import PreTrainedTokenizer, Trie class ModelSetupEmbeddingMixin(metaclass=ABCMeta): def __init__(self): super().__init__() - def _remove_added_embeddings_from_tokenizer( - self, - tokenizer: PreTrainedTokenizer, - ): - if tokenizer: - added_tokens = list(filter(lambda item: not item[1].special, tokenizer._added_tokens_decoder.items())) - for key, added_token in added_tokens: - tokenizer._added_tokens_decoder.pop(key) - tokenizer._added_tokens_encoder.pop(added_token.content) - tokenizer.tokens_trie = Trie() - tokenizer._update_trie() - def _create_new_embedding( self, model: BaseModel, diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index f7c97c132..605994e57 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -118,10 +118,7 @@ def start(self): if self.config.secrets.huggingface_token != "": self.callbacks.on_update_status("logging into Hugging Face") with contextlib.suppress(ConnectionError): - huggingface_hub.login( - token = self.config.secrets.huggingface_token, - new_session = False, - ) + huggingface_hub.login(token=self.config.secrets.huggingface_token) self.callbacks.on_update_status("loading the model") diff --git a/modules/ui/ConceptWindow.py b/modules/ui/ConceptWindow.py index f58879d5f..824ff43e4 100644 --- a/modules/ui/ConceptWindow.py +++ b/modules/ui/ConceptWindow.py @@ -575,7 +575,8 @@ def get_concept_path(path: str) -> str | None: def __download_dataset(self): try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + if self.train_config.secrets.huggingface_token != "": + huggingface_hub.login(token=self.train_config.secrets.huggingface_token) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() diff --git a/modules/util/thread_safety.py b/modules/util/thread_safety.py deleted file mode 100644 index 365076c2a..000000000 --- a/modules/util/thread_safety.py +++ /dev/null @@ -1,43 +0,0 @@ -import functools -import threading - -import torch - -_THREAD_SAFE_FORWARD_ATTR = "_thread_safe_forward_lock" - - -def apply_thread_safe_forward(model: torch.nn.Module) -> None: - """ - Wrap ``model.forward()`` with a per-instance ``threading.Lock`` to - serialize concurrent calls. - - This is a workaround for a thread-safety bug in the transformers library's - ``check_model_inputs`` decorator, which monkey-patches child module - ``.forward()`` methods during execution and is not safe for concurrent use - from multiple dataloader threads. - - See: https://github.com/huggingface/transformers/issues/42673 - Fix: https://github.com/huggingface/transformers/pull/43765 (v5 only) - - This patch can be removed when upgrading to transformers v5+. - - The lock is per-model-instance so different model instances do not block - each other. The function is idempotent: calling it twice on the same model - is a no-op. - - Args: - model: The ``nn.Module`` whose ``forward()`` should be made thread-safe. - """ - if hasattr(model, _THREAD_SAFE_FORWARD_ATTR): - return - - lock = threading.Lock() - original_forward = model.forward - - @functools.wraps(original_forward) - def locked_forward(*args, **kwargs): - with lock: - return original_forward(*args, **kwargs) - - model.forward = locked_forward - setattr(model, _THREAD_SAFE_FORWARD_ATTR, lock) diff --git a/requirements-global.txt b/requirements-global.txt index 91014c6ff..ba3d9f476 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -5,7 +5,7 @@ pillow==12.2.0 imagesize==1.4.1 #for concept statistics tqdm==4.67.1 PyYAML==6.0.2 -huggingface-hub==0.34.4 +huggingface-hub==1.16.1 scipy==1.15.3 matplotlib==3.10.3 av==16.1.0 @@ -22,7 +22,7 @@ tensorboard==2.20.0 # diffusion models -e git+https://github.com/huggingface/diffusers.git@0f1abc4#egg=diffusers gguf==0.17.1 -transformers==4.57.6 +transformers==5.5.4 # pinned below 5.6, see https://github.com/Nerogar/OneTrainer/pull/1506 sentencepiece==0.2.1 # transitive dependency of transformers for tokenizer loading omegaconf==2.3.0 # needed to load stable diffusion from single ckpt files invisible-watermark==0.2.0 # needed for the SDXL pipeline diff --git a/run-cmd.sh b/run-cmd.sh index 2b1c9d14f..8fc77cb02 100755 --- a/run-cmd.sh +++ b/run-cmd.sh @@ -2,11 +2,6 @@ set -e -# Xet is buggy. Disabled by default unless already defined - https://github.com/Nerogar/OneTrainer/issues/949 -if [[ -z "${HF_HUB_DISABLE_XET+x}" ]]; then - export HF_HUB_DISABLE_XET=1 -fi - source "${BASH_SOURCE[0]%/*}/lib.include.sh" # Fetch and validate the name of the target script. diff --git a/start-ui.bat b/start-ui.bat index 5edc8e6c4..05881410b 100644 --- a/start-ui.bat +++ b/start-ui.bat @@ -35,14 +35,6 @@ set PYTHON="%VENV_DIR%\Scripts\python.exe" -X utf8 if defined PROFILE (set PYTHON=%PYTHON% -m scalene --off --cpu --gpu --profile-all --no-browser) echo Using Python %PYTHON% -REM Disable HF_HUB_DISABLE_XET, buggy; default disables Xet (set to 0 to enable) - https://github.com/Nerogar/OneTrainer/issues/949 -if not defined HF_HUB_DISABLE_XET ( - set "HF_HUB_DISABLE_XET=1" -) -echo HF_HUB_DISABLE_XET=%HF_HUB_DISABLE_XET% -echo. -echo NOTE: Xet disabled, to enable it set as 0 before launch - :check_python_version echo Checking Python version... %PYTHON% --version diff --git a/start-ui.sh b/start-ui.sh index 2f0eecc0d..b2960c262 100755 --- a/start-ui.sh +++ b/start-ui.sh @@ -4,11 +4,6 @@ set -e source "${BASH_SOURCE[0]%/*}/lib.include.sh" -# Xet is buggy. Disabled by default unless already defined - https://github.com/Nerogar/OneTrainer/issues/949 -if [[ -z "${HF_HUB_DISABLE_XET+x}" ]]; then - export HF_HUB_DISABLE_XET=1 -fi - prepare_runtime_environment run_python_in_active_env "scripts/train_ui.py" "$@"