Update app.py

app.py

@@ -1,12 +1,9 @@
 import os
 import subprocess
 from typing import Union
-from huggingface_hub import whoami
-is_spaces = True if os.environ.get("SPACE_ID") else False
+from huggingface_hub import whoami, HfApi
 
-if is_spaces:
-    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-    import spaces
+is_spaces = True if os.environ.get("SPACE_ID") else False
     
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 import sys
@@ -22,7 +19,6 @@ import gradio as gr
 from PIL import Image
 import torch
 import uuid
-import os
 import shutil
 import json
 import yaml
@@ -38,16 +34,11 @@ 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)
@@ -203,11 +194,13 @@ def start_training(
     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  # L40S has enough VRAM
+    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)
@@ -215,18 +208,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 = 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
     config["config"]["process"][0]["save"]["hf_token"] = HF_TOKEN
+    config["config"]["process"][0]["model"]["name_or_path"] = "black-forest-labs/FLUX.1-dev"
+    config["config"]["process"][0]["model"]["assistant_lora_path"] = "ostris/FLUX.1-dev-training-adapter"
+    config["config"]["process"][0]["sample"]["sample_steps"] = 28  # dev 모델의 기본 스텝
+   
 
 
-    
     if concept_sentence:
         config["config"]["process"][0]["trigger_word"] = concept_sentence
     
@@ -244,11 +234,6 @@ def start_training(
     else:
         config["config"]["process"][0]["train"]["disable_sampling"] = True
         
-    if(which_model == "[schnell] (4 step fast model)"):
-    # schnell 관련 조건문을 dev로 변경
-        config["config"]["process"][0]["model"]["name_or_path"] = "black-forest-labs/FLUX.1-dev"
-        config["config"]["process"][0]["model"]["assistant_lora_path"] = "ostris/FLUX.1-dev-training-adapter"
-        config["config"]["process"][0]["sample"]["sample_steps"] = 28  # dev 모델의 기본 스텝
 
 
     if(use_more_advanced_options):
@@ -375,12 +360,15 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         placeholder="uncommon word like p3rs0n or trtcrd, or sentence like 'in the style of CNSTLL'",
                         interactive=True,
                     )
+# model_warning 변수 추가
+                model_warning = gr.Markdown(visible=False)
 
                 which_model = gr.Radio(
                     ["[dev] (high quality model)"], 
                     label="Base model", 
                     value="[dev] (high quality model)"
                 )
+     
                     
             with gr.Group(visible=True) as image_upload:
                 with gr.Row():