Update app.py

app.py

@@ -157,6 +157,7 @@ def recursive_update(d, u):
             d[k] = v
     return d
 
+# start_training 함수 수정 부분
 def start_training(
     lora_name,
     concept_sentence,
@@ -173,29 +174,19 @@ def start_training(
     profile: Union[gr.OAuthProfile, None],
     oauth_token: Union[gr.OAuthToken, None],
 ):
-    
     if not lora_name:
         raise gr.Error("You forgot to insert your LoRA name! This name has to be unique.")
     
-    if not is_spaces:
-        try:
-            if whoami()["auth"]["accessToken"]["role"] == "write" or "repo.write" in whoami()["auth"]["accessToken"]["fineGrained"]["scoped"][0]["permissions"]:
-                gr.Info(f"Starting training locally {whoami()['name']}. Your LoRA will be available locally and in Hugging Face after it finishes.")
-            else:
-                raise gr.Error(f"You logged in to Hugging Face with not enough permissions, you need a token that allows writing to your profile.")
-        except: 
-            raise gr.Error(f"You logged in to Hugging Face with not enough permissions, you need a token that allows writing to your profile.")
-            
     print("Started training")
     slugged_lora_name = slugify(lora_name)
 
     # 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:
+    with open("train_lora_flux_24gb.yaml", "r") as f:
         config = yaml.safe_load(f)
 
     # Update the config with user inputs
     config["config"]["name"] = slugged_lora_name
-    config["config"]["process"][0]["model"]["low_vram"] = True
+    config["config"]["process"][0]["model"]["low_vram"] = False  # L40S has enough VRAM
     config["config"]["process"][0]["train"]["skip_first_sample"] = True
     config["config"]["process"][0]["train"]["steps"] = int(steps)
     config["config"]["process"][0]["train"]["lr"] = float(lr)
@@ -203,12 +194,15 @@ 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
+    
     try:
-        username = whoami()["name"] if not is_spaces else profile.username
+        username = profile.username
     except:
         raise gr.Error("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
+    
     if concept_sentence:
         config["config"]["process"][0]["trigger_word"] = concept_sentence
     
@@ -237,39 +231,20 @@ def start_training(
         print(config)
     
     # Save the updated config
-    # generate a random name for the config
     random_config_name = str(uuid.uuid4())
     os.makedirs("tmp", exist_ok=True)
     config_path = f"tmp/{random_config_name}-{slugged_lora_name}.yaml"
     with open(config_path, "w") as f:
         yaml.dump(config, f)
-    if is_spaces:
-        # copy config to dataset_folder as config.yaml
-        shutil.copy(config_path, dataset_folder + "/config.yaml")
-        # get location of this script
-        script_location = os.path.dirname(os.path.abspath(__file__))
-        # copy script.py from current directory to dataset_folder
-        shutil.copy(script_location + "/script.py", dataset_folder)
-        # copy requirements.autotrain to dataset_folder as requirements.txt
-        shutil.copy(script_location + "/requirements.autotrain", dataset_folder + "/requirements.txt")
-        # command to run autotrain spacerunner
-        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"
-        outcome = subprocess.run(cmd.split())
-        if outcome.returncode == 0:
-            return f"""# Your training has started. 
-    ## - Training Status: <a href='https://huggingface.co/spaces/{profile.username}/autotrain-{slugged_lora_name}?logs=container'>{profile.username}/autotrain-{slugged_lora_name}</a> <small>(in the logs tab)</small>
-    ## - Model page: <a href='https://huggingface.co/{profile.username}/{slugged_lora_name}'>{profile.username}/{slugged_lora_name}</a> <small>(will be available when training finishes)</small>"""
-        else:
-            print("Error: ", outcome.stderr)
-            raise gr.Error("Something went wrong. Make sure the name of your LoRA is unique and try again")
-    else:
-        # run the job locally
-        job = get_job(config_path)
-        job.run()
-        job.cleanup()
 
-    return f"Training completed successfully. Model saved as {slugged_lora_name}"
+    # 직접 로컬 GPU에서 학습 실행
+    from toolkit.job import get_job
+    job = get_job(config_path)
+    job.run()
+    job.cleanup()
+
+    return f"""# Training completed successfully! 
+    ## Your model is available at: <a href='https://huggingface.co/{username}/{slugged_lora_name}'>{username}/{slugged_lora_name}</a>"""
 
 def swap_visibilty(profile: Union[gr.OAuthProfile, None]):
     if is_spaces: