Update app.py

app.py

@@ -194,6 +194,14 @@ def start_training(
     print("Started training")
     slugged_lora_name = slugify(lora_name)
 
+    # Create model repo
+    model_repo_id = f"{profile.username}/{slugged_lora_name}"
+    try:
+        api.create_repo(repo_id=model_repo_id, repo_type="model", private=True)
+        print(f"Created model repo {model_repo_id}")
+    except Exception as e:
+        print(f"Model repo might already exist: {e}")
+    
     # Load the default config
     with open("train_lora_flux_24gb.yaml" if is_spaces else "ai-toolkit/config/examples/train_lora_flux_24gb.yaml", "r") as f:
         config = yaml.safe_load(f)
@@ -208,6 +216,7 @@ def start_training(
     config["config"]["process"][0]["network"]["linear_alpha"] = int(rank)
     config["config"]["process"][0]["datasets"][0]["folder_path"] = dataset_folder
     config["config"]["process"][0]["save"]["push_to_hub"] = True
+    config["config"]["process"][0]["save"]["hf_repo_id"] = model_repo_id
     try:
         username = whoami()["name"] if not is_spaces else profile.username
     except:
@@ -307,12 +316,12 @@ def start_training(
             commit_message="Pushing training files from controller"
         )
 
-        model_repo_id = f"{profile.username}/{slugged_lora_name}"
-        try:
-            api.create_repo(repo_id=model_repo_id, repo_type="model", private=True)
-            print(f"Created model repo {model_repo_id}")
-        except Exception as e:
-            print(f"Model repo might already exist: {e}")
+#        model_repo_id = f"{profile.username}/{slugged_lora_name}"
+#        try:
+#            api.create_repo(repo_id=model_repo_id, repo_type="model", private=True)
+#            print(f"Created model repo {model_repo_id}")
+#        except Exception as e:
+#            print(f"Model repo might already exist: {e}")
 
         
         cmd = f"autotrain spacerunner --project-name {slugged_lora_name} --script-path {dataset_folder}"