Merge pull request #9 from neph1/update-v0.7.0

add rename lora script

Author: neph1

app.py

@@ -2,10 +2,9 @@
 
 import gradio as gr
 
-from config import global_config
 from tabs.general_tab import GeneralTab
 from tabs.prepare_tab import PrepareDatasetTab
-from tabs.tab import Tab
+from tabs.tool_tab import ToolTab
 from tabs.training_tab import TrainingTab
 from tabs.training_tab_legacy import LegacyTrainingTab
 
@@ -22,28 +21,32 @@ def __init__(self):
         self.setup_views()
 
     def setup_views(self):
-        with gr.Blocks() as demo:
+        with gr.Blocks() as app:
             gr.Markdown("### finetrainers ui")
 
 
             with gr.Tab("General Settings"):
                 self.tabs['general'] = GeneralTab("General Settings", os.path.join(self.configs_path, "editor.yaml"))
             runtime_tab = gr.Tab("Trainer Settings")
+            tools_tab = gr.Tab("Lora Tools")
+            
             prepare_tab = gr.Tab("Prepare dataset (Legacy)")
             runtime_tab_legacy = gr.Tab("Legacy Training Settings")
 
             with runtime_tab:
                 self.tabs['runtime'] = TrainingTab("Trainer Settings", os.path.join(self.configs_path, "config_template.yaml"), allow_load=True)
 
-            
+            with tools_tab:
+                self.tabs['tools'] = ToolTab()
+
             with prepare_tab:
                 self.tabs['prepare'] = PrepareDatasetTab("Prepare dataset (Legacy)", os.path.join(self.configs_path, "prepare_template.yaml"), allow_load=True)
 
             with runtime_tab_legacy:
                 self.tabs['runtime'] = LegacyTrainingTab("Legacy CogvideoX Settings", os.path.join(self.configs_path, "config_template_legacy.yaml"), allow_load=True)
 
 
-        demo.launch()
+        app.launch()
 
 
 if __name__ == "__main__":

requirements.txt

@@ -1 +1,2 @@
-gradio
+gradio
+torch>=2.4.1

scripts/__init__.py

@@ -0,0 +1,28 @@
+import os
+import torch
+from safetensors.torch import load_file, save_file, safe_open
+
+def load_state_dict(file_name, dtype):
+    if os.path.splitext(file_name)[1] == ".safetensors":
+        sd = load_file(file_name)
+        with safe_open(file_name, framework="pt") as f:
+            metadata = f.metadata()
+    else:
+        sd = torch.load(file_name, map_location="cpu")
+        metadata = None
+
+    for key in list(sd.keys()):
+        if type(sd[key]) == torch.Tensor:
+            sd[key] = sd[key].to(dtype)
+
+    return sd, metadata
+
+def save_to_file(file_name, state_dict, dtype, metadata):
+    if dtype is not None:
+        for key in list(state_dict.keys()):
+            if type(state_dict[key]) == torch.Tensor:
+                state_dict[key] = state_dict[key].to(dtype) 
+    if os.path.splitext(file_name)[1] == ".safetensors":
+        save_file(state_dict, file_name, metadata)
+    else:
+        torch.save(state_dict, file_name)

scripts/common_io.py

@@ -0,0 +1,29 @@
+import argparse
+import torch
+
+from scripts.common_io import load_state_dict, save_to_file
+
+
+def rename_keys(file, outfile: str)-> bool:
+    sd, metadata = load_state_dict(file, torch.float32)
+
+    keys_to_normalize = [key for key in sd.keys()]
+    values_to_normalize = [sd[key].to(torch.float32) for key in keys_to_normalize]
+    new_sd = dict()
+    for key, value in zip(keys_to_normalize, values_to_normalize):
+        new_sd[key.replace("transformer.", "")] = value
+        
+    save_to_file(outfile, new_sd, torch.float16, metadata)
+    return True
+
+def setup_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("file", type=str, help="lora model to modify")
+    parser.add_argument("-o", "--outfile", type=str, help="file to save to")
+    return parser
+
+
+if __name__ == "__main__":
+    parser = setup_parser()
+    args = parser.parse_args()
+    rename_keys(args.file, args.outfile)

scripts/rename_keys.py

@@ -0,0 +1,22 @@
+import gradio as gr
+
+from scripts import rename_keys
+
+class ToolTab:
+
+    def __init__(self):
+        gr.Markdown("Lora Tools")
+        self._build_rename()
+
+    def _build_rename(self):
+        gr.Markdown("Rename Transformer Keys")
+        rename_in_file = gr.Textbox(value="", label="Lora to rename (full path)", lines=1)
+        rename_out_file = gr.Textbox(value="", label="New name (full path)", lines=1)
+        rename_result = gr.Textbox(value="", label="Result", lines=1)
+        rename_button = gr.Button("Rename", key='rename')
+        rename_button.click(self.rename_file, inputs=[rename_in_file, rename_out_file], outputs=[rename_result])
+
+    def rename_file(self, lora_name, new_name):
+        result = rename_keys.rename_keys(lora_name, new_name)
+        if result:
+            return "Renamed successfully"