diff --git a/modules/modelSaver/chroma/ChromaModelSaver.py b/modules/modelSaver/chroma/ChromaModelSaver.py index f8129558d..f76a90750 100644 --- a/modules/modelSaver/chroma/ChromaModelSaver.py +++ b/modules/modelSaver/chroma/ChromaModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -29,18 +28,7 @@ def __save_diffusers( # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer = pipeline.tokenizer - tokenizer.__deepcopy__ = lambda memo: tokenizer - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype) text_encoder = save_pipeline.text_encoder if text_encoder is not None: diff --git a/modules/modelSaver/ernie/ErnieModelSaver.py b/modules/modelSaver/ernie/ErnieModelSaver.py index 65ff388d3..a99bd76bf 100644 --- a/modules/modelSaver/ernie/ErnieModelSaver.py +++ b/modules/modelSaver/ernie/ErnieModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -25,16 +24,7 @@ def __save_diffusers( ): pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: - tokenizer = pipeline.tokenizer - tokenizer.__deepcopy__ = lambda memo: tokenizer - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype) os.makedirs(Path(destination).absolute(), exist_ok=True) save_pipeline.save_pretrained(destination) diff --git a/modules/modelSaver/flux/FluxModelSaver.py b/modules/modelSaver/flux/FluxModelSaver.py index dc632f2ae..f98b4cf46 100644 --- a/modules/modelSaver/flux/FluxModelSaver.py +++ b/modules/modelSaver/flux/FluxModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -29,18 +28,7 @@ def __save_diffusers( # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer_2 = pipeline.tokenizer_2 - tokenizer_2.__deepcopy__ = lambda memo: tokenizer_2 - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer_2, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype, tokenizer_attrs=("tokenizer_2",)) text_encoder_2 = save_pipeline.text_encoder_2 if text_encoder_2 is not None: diff --git a/modules/modelSaver/flux2/Flux2ModelSaver.py b/modules/modelSaver/flux2/Flux2ModelSaver.py index e2976244e..b12d06c86 100644 --- a/modules/modelSaver/flux2/Flux2ModelSaver.py +++ b/modules/modelSaver/flux2/Flux2ModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -27,18 +26,7 @@ def __save_diffusers( # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: #TODO necessary? - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer = pipeline.tokenizer - tokenizer.__deepcopy__ = lambda memo: tokenizer - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype) os.makedirs(Path(destination).absolute(), exist_ok=True) save_pipeline.save_pretrained(destination) diff --git a/modules/modelSaver/hidream/HiDreamModelSaver.py b/modules/modelSaver/hidream/HiDreamModelSaver.py index 4bfe56ebc..5665d73ed 100644 --- a/modules/modelSaver/hidream/HiDreamModelSaver.py +++ b/modules/modelSaver/hidream/HiDreamModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -26,21 +25,7 @@ def __save_diffusers( # 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.to("cpu") - if dtype is not None: - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer_3 = pipeline.tokenizer_3 - tokenizer_3.__deepcopy__ = lambda memo: tokenizer_3 - tokenizer_4 = pipeline.tokenizer_4 - tokenizer_4.__deepcopy__ = lambda memo: tokenizer_4 - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer_3, '__deepcopy__') - delattr(tokenizer_4, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype, tokenizer_attrs=("tokenizer_3", "tokenizer_4")) os.makedirs(Path(destination).absolute(), exist_ok=True) save_pipeline.save_pretrained(destination) diff --git a/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py b/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py index 866d68617..286c61377 100644 --- a/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py +++ b/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -29,18 +28,7 @@ def __save_diffusers( # 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.to("cpu") - if dtype is not None: - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer_1 = pipeline.tokenizer - tokenizer_1.__deepcopy__ = lambda memo: tokenizer_1 - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer_1, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype) text_encoder_1 = save_pipeline.text_encoder if text_encoder_1 is not None: diff --git a/modules/modelSaver/mixin/DtypeModelSaverMixin.py b/modules/modelSaver/mixin/DtypeModelSaverMixin.py index 821477565..5a2d1c835 100644 --- a/modules/modelSaver/mixin/DtypeModelSaverMixin.py +++ b/modules/modelSaver/mixin/DtypeModelSaverMixin.py @@ -42,6 +42,29 @@ def _convert_state_dict_to_contiguous( else: state_dict[key] = value.contiguous() + def _copy_pipeline_to_dtype( + self, + pipeline, + dtype: torch.dtype | None, + tokenizer_attrs: tuple[str, ...] = ("tokenizer",), + ): + if dtype is None: + return pipeline + + # replace the tokenizers' __deepcopy__ before calling deepcopy, to prevent a copy being made. + # the tokenizers try to reload from the file system otherwise + tokenizers = [getattr(pipeline, attr) for attr in tokenizer_attrs] + for tokenizer in tokenizers: + tokenizer.__deepcopy__ = lambda memo, tokenizer=tokenizer: tokenizer + + save_pipeline = copy.deepcopy(pipeline) + save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) + + for tokenizer in tokenizers: + delattr(tokenizer, '__deepcopy__') + + return save_pipeline + def __calculate_safetensors_hash( self, state_dict: dict[str, Tensor] | None = None, diff --git a/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py b/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py index 83e8af4a9..016534229 100644 --- a/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py +++ b/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -28,18 +27,7 @@ def __save_diffusers( pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer = pipeline.tokenizer - tokenizer.__deepcopy__ = lambda memo: tokenizer - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype) os.makedirs(Path(destination).absolute(), exist_ok=True) save_pipeline.save_pretrained(destination) diff --git a/modules/modelSaver/qwen/QwenModelSaver.py b/modules/modelSaver/qwen/QwenModelSaver.py index cee262a8d..80f630bf6 100644 --- a/modules/modelSaver/qwen/QwenModelSaver.py +++ b/modules/modelSaver/qwen/QwenModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -26,19 +25,7 @@ def __save_diffusers( # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: - #TODO is this code necessary for all models? in that case, share code - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer = pipeline.tokenizer - tokenizer.__deepcopy__ = lambda memo: tokenizer - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype) os.makedirs(Path(destination).absolute(), exist_ok=True) save_pipeline.save_pretrained(destination) diff --git a/modules/modelSaver/sana/SanaModelSaver.py b/modules/modelSaver/sana/SanaModelSaver.py index 1dede8c8f..575dfe789 100644 --- a/modules/modelSaver/sana/SanaModelSaver.py +++ b/modules/modelSaver/sana/SanaModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -25,18 +24,7 @@ def __save_diffusers( pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer = pipeline.tokenizer - tokenizer.__deepcopy__ = lambda memo: tokenizer - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype) os.makedirs(Path(destination).absolute(), exist_ok=True) save_pipeline.save_pretrained(destination) diff --git a/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py b/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py index bcfabab67..abaa85178 100644 --- a/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py +++ b/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -29,18 +28,7 @@ def __save_diffusers( # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer_3 = pipeline.tokenizer_3 - tokenizer_3.__deepcopy__ = lambda memo: tokenizer_3 - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer_3, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype, tokenizer_attrs=("tokenizer_3",)) text_encoder_3 = save_pipeline.text_encoder_3 if text_encoder_3 is not None: diff --git a/modules/modelSaver/zImage/ZImageModelSaver.py b/modules/modelSaver/zImage/ZImageModelSaver.py index 1d8f85495..7a46f16dd 100644 --- a/modules/modelSaver/zImage/ZImageModelSaver.py +++ b/modules/modelSaver/zImage/ZImageModelSaver.py @@ -1,4 +1,3 @@ -import copy import os.path from pathlib import Path @@ -26,19 +25,7 @@ def __save_diffusers( # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. pipeline = model.create_pipeline() pipeline.to("cpu") - if dtype is not None: - #TODO is this code necessary for all models? in that case, share code - # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. - # the tokenizer tries to reload from the file system otherwise - tokenizer = pipeline.tokenizer - tokenizer.__deepcopy__ = lambda memo: tokenizer - - save_pipeline = copy.deepcopy(pipeline) - save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) - - delattr(tokenizer, '__deepcopy__') - else: - save_pipeline = pipeline + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype) os.makedirs(Path(destination).absolute(), exist_ok=True) save_pipeline.save_pretrained(destination)