Update flux_train.py

flux_train.py

@@ -2,14 +2,28 @@
 
 import os
 from collections import OrderedDict
+from huggingface_hub import whoami
 import sys
 sys.path.append("/app/ai-toolkit")  # Tell Python to look here
 
+from toolkit.job import run_job
 from toolkit.job import run_job
 
+def update_config_push_to_hub(config, push_to_hub: bool, slugged_lora_name: str):
+    config["config"]["process"][0]["save"]["push_to_hub"] = push_to_hub
 
+    if push_to_hub:
+        try:
+            username = whoami()["name"]
+        except Exception:
+            raise RuntimeError(
+                "Error trying to retrieve your username. Are you sure you are logged in with Hugging Face?"
+            )
+        config["config"]["process"][0]["save"]["hf_repo_id"] = f"{username}/{slugged_lora_name}"
+        config["config"]["process"][0]["save"]["hf_private"] = True
 
-def build_job(concept="ohamlab style", training_path="/tmp/data", lora_name="ohami_filter_autorun"):
+def build_job(concept="ohamlab style", training_path="/tmp/data", lora_name="ohami_filter_autorun", push_to_hub=False):
+    slugged_lora_name = lora_name.lower().replace(" ", "_")
     job = OrderedDict([
         ('job', 'extension'),
         ('config', OrderedDict([
@@ -27,7 +41,8 @@ def build_job(concept="ohamlab style", training_path="/tmp/data", lora_name="oha
                     ('save', OrderedDict([
                         ('dtype', 'float16'),
                         ('save_every', 250),
-                        ('max_step_saves_to_keep', 4)
+                        ('max_step_saves_to_keep', 4),
+                        # push_to_hub keys added later by update_config_push_to_hub()
                     ])),
                     ('datasets', [
                         OrderedDict([
@@ -81,4 +96,11 @@ def build_job(concept="ohamlab style", training_path="/tmp/data", lora_name="oha
             ('version', '1.0')
         ]))
     ])
+
+    # Add push to Hub config if requested
+    update_config_push_to_hub(job, push_to_hub, slugged_lora_name)
     return job
+
+def run_training(concept="ohamlab style", training_path="/tmp/data", lora_name="ohami_filter_autorun", push_to_hub=False):
+    job = build_job(concept, training_path, lora_name, push_to_hub)
+    run_job(job)