Update app.py

app.py

@@ -2,15 +2,9 @@ import os
 import subprocess
 from typing import Union
 from huggingface_hub import whoami
+from huggingface_hub import HfApi
 is_spaces = True if os.environ.get("SPACE_ID") else False
 
-#from huggingface_hub import login
-#hf_token = os.getenv("HF_TOKEN")
-#if hf_token:
-#    login(token=hf_token)
-#else:
-#    raise ValueError("HF_TOKEN environment variable not found!")
-
 if is_spaces:
     subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
     import spaces
@@ -260,8 +254,27 @@ def start_training(
         # copy requirements.autotrain to dataset_folder as requirements.txt
         shutil.copy(script_location + "/requirements.autotrain", dataset_folder + "/requirements.txt")
         # command to run autotrain spacerunner
+        # --- Pre-create subspace and inject token ---
+        hf_token = oauth_token.token
+        api = HfApi(token=hf_token)
+        new_space_id = f"{profile.username}/autotrain-{slugged_lora_name}"
+
+        try:
+            # Try to create the space (if it doesn't exist)
+            api.create_repo(repo_id=new_space_id, repo_type="space", space_sdk="docker", private=True)
+            print(f"Created space {new_space_id}")
+        except Exception as e:
+            print(f"Space creation may have failed or already exists: {e}")
+
+        try:
+            # Inject HF_TOKEN into the new space
+            api.add_space_secret(repo_id=new_space_id, key="HF_TOKEN", value=hf_token)
+            print(f"Added HF_TOKEN secret to {new_space_id}")
+        except Exception as e:
+            print(f"Failed to add HF_TOKEN secret to new space: {e}")
+            
         cmd = f"autotrain spacerunner --project-name {slugged_lora_name} --script-path {dataset_folder}"
-        cmd += f" --username {profile.username} --token {oauth_token.token} --backend spaces-l4x1"
+        cmd += f" --username {profile.username} --token {hf_token} --backend spaces-l4x1"
         outcome = subprocess.run(cmd.split())
         if outcome.returncode == 0:
             return f"""# Your training has started.