Update app.py

app.py

@@ -204,8 +204,25 @@ def start_training(
     use_more_advanced_options,
     more_advanced_options,
 ):
-    if not lora_name:
-        raise gr.Error("You forgot to insert your LoRA name! This name has to be unique.")
+    # 환경 변수로 타임아웃 설정
+    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)
     
     try:
         username = whoami()["name"]
@@ -215,71 +232,25 @@ def start_training(
     print("Started training")
     slugged_lora_name = slugify(lora_name)
 
-    try:
-        from toolkit.job import get_job
-    except ImportError:
-        raise gr.Error("Failed to import toolkit. Please check if ai-toolkit is properly installed.")
-        
-    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
-    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
+    # dev 모델 설정
     config["config"]["process"][0]["model"]["name_or_path"] = "black-forest-labs/FLUX.1-dev"
-    config["config"]["process"][0]["model"]["assistant_lora_path"] = None  # training adapter 없이 시도
-    config["config"]["process"][0]["sample"]["sample_steps"] = 28  # dev 모델의 기본 스텝
-   
-    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"]["sample_steps"] = 28
-        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
-        
+    config["config"]["process"][0]["model"]["assistant_lora_path"] = None  # adapter 없이 설정
+    config["config"]["process"][0]["sample"]["sample_steps"] = 28
 
+    # 나머지 설정은 동일...
 
-    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)
-    
-    # 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)
-    job.run()
-    job.cleanup()
+    try:
+        # 직접 로컬 GPU에서 학습 실행
+        from toolkit.job import get_job
+        job = get_job(config_path)
+        job.run()
+        job.cleanup()
+    except Exception as e:
+        raise gr.Error(f"Training failed: {str(e)}")
 
     return f"""# Training completed successfully! 
     ## Your model is available at: <a href='https://huggingface.co/{username}/{slugged_lora_name}'>{username}/{slugged_lora_name}</a>"""