Model saver code deduplication

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:

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)

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:

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)

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)

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:

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,

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)

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)

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)

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:

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)