add tools, update login page

src/autotrain/cli/autotrain.py

@@ -12,6 +12,7 @@
 from .run_tabular import RunAutoTrainTabularCommand
 from .run_text_classification import RunAutoTrainTextClassificationCommand
 from .run_token_classification import RunAutoTrainTokenClassificationCommand
+from .run_tools import RunAutoTrainToolsCommand
 
 
 def main():
@@ -35,6 +36,7 @@ def main():
     RunAutoTrainSpaceRunnerCommand.register_subcommand(commands_parser)
     RunAutoTrainSeq2SeqCommand.register_subcommand(commands_parser)
     RunAutoTrainTokenClassificationCommand.register_subcommand(commands_parser)
+    RunAutoTrainToolsCommand.register_subcommand(commands_parser)
 
     args = parser.parse_args()
 

src/autotrain/cli/run_tools.py

@@ -0,0 +1,99 @@
+from argparse import ArgumentParser
+
+from . import BaseAutoTrainCommand
+
+
+def run_tools_command_factory(args):
+    return RunAutoTrainToolsCommand(args)
+
+
+class RunAutoTrainToolsCommand(BaseAutoTrainCommand):
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        run_app_parser = parser.add_parser("tools", help="Run AutoTrain tools")
+        subparsers = run_app_parser.add_subparsers(title="tools", dest="tool_name")
+
+        merge_llm_parser = subparsers.add_parser(
+            "merge-llm-adapter",
+            help="Merge LLM Adapter tool",
+        )
+        merge_llm_parser.add_argument(
+            "--base-model-path",
+            type=str,
+            help="Base model path",
+        )
+        merge_llm_parser.add_argument(
+            "--adapter-path",
+            type=str,
+            help="Adapter path",
+        )
+        merge_llm_parser.add_argument(
+            "--token",
+            type=str,
+            help="Token",
+            default=None,
+            required=False,
+        )
+        merge_llm_parser.add_argument(
+            "--pad-to-multiple-of",
+            type=int,
+            help="Pad to multiple of",
+            default=None,
+            required=False,
+        )
+        merge_llm_parser.add_argument(
+            "--output-folder",
+            type=str,
+            help="Output folder",
+            required=False,
+            default=None,
+        )
+        merge_llm_parser.add_argument(
+            "--push-to-hub",
+            action="store_true",
+            help="Push to Hugging Face Hub",
+            required=False,
+        )
+        merge_llm_parser.set_defaults(func=run_tools_command_factory, merge_llm_adapter=True)
+
+        convert_to_kohya_parser = subparsers.add_parser("convert_to_kohya", help="Convert to Kohya tool")
+        convert_to_kohya_parser.add_argument(
+            "--input-path",
+            type=str,
+            help="Input path",
+        )
+        convert_to_kohya_parser.add_argument(
+            "--output-path",
+            type=str,
+            help="Output path",
+        )
+        convert_to_kohya_parser.set_defaults(func=run_tools_command_factory, convert_to_kohya=True)
+
+    def __init__(self, args):
+        self.args = args
+
+    def run(self):
+        if getattr(self.args, "merge_llm_adapter", False):
+            self.run_merge_llm_adapter()
+        if getattr(self.args, "convert_to_kohya", False):
+            self.run_convert_to_kohya()
+
+    def run_merge_llm_adapter(self):
+        from autotrain.tools.merge_adapter import merge_llm_adapter
+
+        merge_llm_adapter(
+            base_model_path=self.args.base_model_path,
+            adapter_path=self.args.adapter_path,
+            token=self.args.token,
+            output_folder=self.args.output_folder,
+            pad_to_multiple_of=self.args.pad_to_multiple_of,
+            push_to_hub=self.args.push_to_hub,
+        )
+
+    def run_convert_to_kohya(self):
+        from autotrain.tools.convert_to_kohya import convert_to_kohya
+
+        convert_to_kohya(
+            input_path=self.args.input_path,
+            output_path=self.args.output_path,
+        )

src/autotrain/templates/login.html

@@ -7,6 +7,13 @@
     <script src="https://cdn.tailwindcss.com"></script>
 </head>
 
+<script>
+    function redirectToLogin() {
+        const uri = '/login/huggingface';
+        window.location.assign(uri + window.location.search);
+    }
+</script>
+
 <body>
     <header class="bg-white-800 text-white p-4">
         <div class="container mx-auto flex justify-between items-center">
@@ -15,15 +22,23 @@
     </header>
 
     <div class="form-container max-w-lg mx-auto mt-10 p-6 shadow-2xl">
-        <h1 class="text-2xl font-bold mb-10">Login</h1>
-        <p class="text-gray-500 text-sm mb-10">Please <a href="/login/huggingface" target="_blank">login</a> to use
+        <p class="text-gray-500 text-xl mb-3 text-center">Please <a href="javascript:void(0);"
+                onclick="redirectToLogin()">login</a> to use
             AutoTrain</p>
-        <div style="text-align: center;">
-            <a href="/login/huggingface" target="_blank">
+        <div class="flex justify-center items-center">
+            <a href="javascript:void(0);" onclick="redirectToLogin()">
                 <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/sign-in-with-huggingface-xl.svg"
                     alt="Login using Hugging Face">
             </a>
         </div>
+        <p class="text-gray-500 text-xs mt-10 text-center">Alternatively, if you face login issues, you can add your
+            Hugging Face Write Token to this space as a secret in space settings. Note: The name of secret must be
+            HF_TOKEN and the value must be your Hugging Face WRITE token! You can find your tokens in user settings.</p>
+        <div class="block text-sm font-normal text-gray-700 text-center"><a href="https://hf.co/docs/autotrain"
+                target="_blank" class="text-blue-600">Docs</a> | <a
+                href="https://github.com/huggingface/autotrain-advanced" target="_blank"
+                class="text-blue-600">GitHub</a>
+        </div>
     </div>
 </body>
 

src/autotrain/tools/convert_to_kohya.py

@@ -0,0 +1,13 @@
+from diffusers.utils import convert_all_state_dict_to_peft, convert_state_dict_to_kohya
+from safetensors.torch import load_file, save_file
+
+from autotrain import logger
+
+
+def convert_to_kohya(input_path, output_path):
+    logger.info(f"Converting Lora state dict from {input_path} to Kohya state dict at {output_path}")
+    lora_state_dict = load_file(input_path)
+    peft_state_dict = convert_all_state_dict_to_peft(lora_state_dict)
+    kohya_state_dict = convert_state_dict_to_kohya(peft_state_dict)
+    save_file(kohya_state_dict, output_path)
+    logger.info(f"Kohya state dict saved at {output_path}")

src/autotrain/tools/merge_adapter.py

@@ -0,0 +1,52 @@
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from autotrain import logger
+from autotrain.trainers.common import ALLOW_REMOTE_CODE
+
+
+def merge_llm_adapter(
+    base_model_path, adapter_path, token, output_folder=None, pad_to_multiple_of=None, push_to_hub=False
+):
+
+    if output_folder is None and push_to_hub is False:
+        raise ValueError("You must specify either --output_folder or --push_to_hub")
+
+    logger.info("Loading adapter...")
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_model_path,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+        trust_remote_code=ALLOW_REMOTE_CODE,
+        token=token,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        adapter_path,
+        trust_remote_code=ALLOW_REMOTE_CODE,
+        token=token,
+    )
+    if pad_to_multiple_of:
+        base_model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=pad_to_multiple_of)
+    else:
+        base_model.resize_token_embeddings(len(tokenizer))
+
+    model = PeftModel.from_pretrained(
+        base_model,
+        adapter_path,
+        token=token,
+    )
+    model = model.merge_and_unload()
+
+    if output_folder is not None:
+        logger.info("Saving target model...")
+        model.save_pretrained(output_folder)
+        tokenizer.save_pretrained(output_folder)
+        logger.info(f"Model saved to {output_folder}")
+
+    if push_to_hub:
+        logger.info("Pushing model to Hugging Face Hub...")
+        model.push_to_hub(adapter_path)
+        tokenizer.push_to_hub(adapter_path)
+        logger.info(f"Model pushed to Hugging Face Hub as {adapter_path}")