Update app.py

app.py

@@ -38,6 +38,27 @@ if not is_spaces:
 MAX_IMAGES = 150
 
 
+import subprocess
+from typing import Union
+from huggingface_hub import whoami, HfApi
+
+# Hugging Face 토큰 설정
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable is not set")
+
+is_spaces = True if os.environ.get("SPACE_ID") else False
+
+if is_spaces:
+    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+    import spaces
+    
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN
+
+# HF API 초기화
+api = HfApi(token=HF_TOKEN)
+
 def load_captioning(uploaded_files, concept_sentence):
     uploaded_images = [file for file in uploaded_files if not file.endswith('.txt')]
     txt_files = [file for file in uploaded_files if file.endswith('.txt')]
@@ -157,7 +178,6 @@ def recursive_update(d, u):
             d[k] = v
     return d
 
-# start_training 함수 수정 부분
 def start_training(
     lora_name,
     concept_sentence,
@@ -171,18 +191,19 @@ def start_training(
     sample_3,
     use_more_advanced_options,
     more_advanced_options,
-    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.")
     
+    try:
+        username = whoami()["name"]
+    except:
+        raise gr.Error("Failed to get username. Please check your HF_TOKEN.")
+        
     print("Started training")
     slugged_lora_name = slugify(lora_name)
 
-    # Load the default config
-    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
@@ -202,6 +223,9 @@ def start_training(
     
     config["config"]["process"][0]["save"]["hf_repo_id"] = f"{username}/{slugged_lora_name}"
     config["config"]["process"][0]["save"]["hf_private"] = True
+    config["config"]["process"][0]["save"]["hf_token"] = HF_TOKEN
+
+
     
     if concept_sentence:
         config["config"]["process"][0]["trigger_word"] = concept_sentence