Upload jobs/process/BaseExtractProcess.py with huggingface_hub

jobs/process/BaseExtractProcess.py

@@ -0,0 +1,86 @@
+import os
+from collections import OrderedDict
+
+from safetensors.torch import save_file
+
+from jobs.process.BaseProcess import BaseProcess
+from toolkit.metadata import get_meta_for_safetensors
+
+from typing import ForwardRef
+
+from toolkit.train_tools import get_torch_dtype
+
+
+class BaseExtractProcess(BaseProcess):
+
+    def __init__(
+            self,
+            process_id: int,
+            job,
+            config: OrderedDict
+    ):
+        super().__init__(process_id, job, config)
+        self.config: OrderedDict
+        self.output_folder: str
+        self.output_filename: str
+        self.output_path: str
+        self.process_id = process_id
+        self.job = job
+        self.config = config
+        self.dtype = self.get_conf('dtype', self.job.dtype)
+        self.torch_dtype = get_torch_dtype(self.dtype)
+        self.extract_unet = self.get_conf('extract_unet', self.job.extract_unet)
+        self.extract_text_encoder = self.get_conf('extract_text_encoder', self.job.extract_text_encoder)
+
+    def run(self):
+        # here instead of init because child init needs to go first
+        self.output_path = self.get_output_path()
+        # implement in child class
+        # be sure to call super().run() first
+        pass
+
+    # you can override this in the child class if you want
+    # call super().get_output_path(prefix="your_prefix_", suffix="_your_suffix") to extend this
+    def get_output_path(self, prefix=None, suffix=None):
+        config_output_path = self.get_conf('output_path', None)
+        config_filename = self.get_conf('filename', None)
+        # replace [name] with name
+
+        if config_output_path is not None:
+            config_output_path = config_output_path.replace('[name]', self.job.name)
+            return config_output_path
+
+        if config_output_path is None and config_filename is not None:
+            # build the output path from the output folder and filename
+            return os.path.join(self.job.output_folder, config_filename)
+
+        # build our own
+
+        if suffix is None:
+            # we will just add process it to the end of the filename if there is more than one process
+            # and no other suffix was given
+            suffix = f"_{self.process_id}" if len(self.config['process']) > 1 else ''
+
+        if prefix is None:
+            prefix = ''
+
+        output_filename = f"{prefix}{self.output_filename}{suffix}"
+
+        return os.path.join(self.job.output_folder, output_filename)
+
+    def save(self, state_dict):
+        # prepare meta
+        save_meta = get_meta_for_safetensors(self.meta, self.job.name)
+
+        # save
+        os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
+
+        for key in list(state_dict.keys()):
+            v = state_dict[key]
+            v = v.detach().clone().to("cpu").to(self.torch_dtype)
+            state_dict[key] = v
+
+        # having issues with meta
+        save_file(state_dict, self.output_path, save_meta)
+
+        print(f"Saved to {self.output_path}")