Update app.py

app.py

@@ -204,25 +204,8 @@ def start_training(
     use_more_advanced_options,
     more_advanced_options,
 ):
-    # 환경 변수로 타임아웃 설정
-    os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "300"  # 5분으로 증가
-    os.environ["REQUESTS_TIMEOUT"] = "300"
-    
-    import requests
-    from huggingface_hub import HfApi
-    from requests.adapters import HTTPAdapter
-    from urllib3.util.retry import Retry
-    
-    # 재시도 전략 설정
-    retry_strategy = Retry(
-        total=5,
-        backoff_factor=1,
-        status_forcelist=[429, 500, 502, 503, 504],
-    )
-    adapter = HTTPAdapter(max_retries=retry_strategy)
-    http = requests.Session()
-    http.mount("https://", adapter)
-    http.mount("http://", adapter)
+    if not lora_name:
+        raise gr.Error("You forgot to insert your LoRA name! This name has to be unique.")
     
     try:
         username = whoami()["name"]
@@ -237,13 +220,51 @@ def start_training(
         config = yaml.safe_load(f)
 
     # dev 모델 설정
+    config["config"]["name"] = slugged_lora_name
     config["config"]["process"][0]["model"]["name_or_path"] = "black-forest-labs/FLUX.1-dev"
     config["config"]["process"][0]["model"]["assistant_lora_path"] = None  # adapter 없이 설정
+    config["config"]["process"][0]["model"]["low_vram"] = False
+    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)
+    config["config"]["process"][0]["network"]["linear"] = int(rank)
+    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"] = f"{username}/{slugged_lora_name}"
+    config["config"]["process"][0]["save"]["hf_private"] = True
+    config["config"]["process"][0]["save"]["hf_token"] = HF_TOKEN
     config["config"]["process"][0]["sample"]["sample_steps"] = 28
 
-    # 나머지 설정은 동일...
-
+    if concept_sentence:
+        config["config"]["process"][0]["trigger_word"] = concept_sentence
+    
+    if sample_1 or sample_2 or sample_3:
+        config["config"]["process"][0]["train"]["disable_sampling"] = False
+        config["config"]["process"][0]["sample"]["sample_every"] = steps
+        config["config"]["process"][0]["sample"]["prompts"] = []
+        if sample_1:
+            config["config"]["process"][0]["sample"]["prompts"].append(sample_1)
+        if sample_2:
+            config["config"]["process"][0]["sample"]["prompts"].append(sample_2)
+        if sample_3:
+            config["config"]["process"][0]["sample"]["prompts"].append(sample_3)
+    else:
+        config["config"]["process"][0]["train"]["disable_sampling"] = True
+
+    if(use_more_advanced_options):
+        more_advanced_options_dict = yaml.safe_load(more_advanced_options)
+        config["config"]["process"][0] = recursive_update(config["config"]["process"][0], more_advanced_options_dict)
+        print(config)
+    
     try:
+        # Save the updated 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)
+
         # 직접 로컬 GPU에서 학습 실행
         from toolkit.job import get_job
         job = get_job(config_path)
@@ -255,6 +276,7 @@ def start_training(
     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 update_pricing(steps):
     try:
         seconds_per_iteration = 7.54