Update flux_train_ui.py

flux_train_ui.py

@@ -2,7 +2,7 @@ import os
 from huggingface_hub import whoami    
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 import sys
-
+import spaces
 # Add the current working directory to the Python path
 sys.path.insert(0, os.getcwd())
 
@@ -140,6 +140,23 @@ def recursive_update(d, u):
             d[k] = v
     return d
 
+
+def get_duration(  lora_name,
+    concept_sentence,
+    steps,
+    lr,
+    rank,
+    model_to_train,
+    low_vram,
+    dataset_folder,
+    sample_1,
+    sample_2,
+    sample_3,
+    use_more_advanced_options,
+    more_advanced_options,):
+    return total_second_length * 60
+
+@spaces.GPU(duration=get_duration)
 def start_training(
     lora_name,
     concept_sentence,