Update app.py

app.py

@@ -35,6 +35,22 @@ def swap_hardware(hf_token, hardware="cpu-basic"):
     headers = { "authorization" : f"Bearer {hf_token}"}
     body = {'flavor': hardware}
     requests.post(hardware_url, json = body, headers=headers)
+
+def swap_sleep_time(hf_token,sleep_time):
+    sleep_time_url = f"https://huggingface.co/api/spaces/{os.environ['SPACE_ID']}/sleeptime"
+    headers = { "authorization" : f"Bearer {hf_token}"}
+    body = {'seconds':sleep_time}
+    requests.post(sleep_time_url,json=body,headers=headers)
+
+def get_sleep_time(hf_token):
+    sleep_time_url = f"https://huggingface.co/api/spaces/{os.environ['SPACE_ID']}"
+    headers = { "authorization" : f"Bearer {hf_token}"}
+    response = requests.get(sleep_time_url,headers=headers)
+    try:
+        gcTimeout = response.json()['runtime']['gcTimeout']
+    except:
+        gcTimeout = None
+    return gcTimeout
     
 def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instance_prompt, class_prompt, max_train_steps, checkpoint_steps):
     
@@ -91,6 +107,10 @@ def main(image_path, b_lora_trained_folder, instance_prompt, class_prompt, train
 
     if instance_prompt == "":
         raise gr.Error("You forgot to specify an instance prompt")
+
+    sleep_time = get_sleep_time(hf_token)
+    if sleep_time:
+        swap_sleep_time(hf_token, -1)
     
     local_dir = "image_to_train"
     # Check if the directory exists and create it if necessary
@@ -111,7 +131,8 @@ def main(image_path, b_lora_trained_folder, instance_prompt, class_prompt, train
     
     your_username = api.whoami(token=hf_token)["name"]
 
-    swap_hardware(hf_token, hardware="cpu-basic")
+    #swap_hardware(hf_token, hardware="cpu-basic")
+    swap_sleep_time(hf_token, 300)
     
     return f"Done, your trained model has been stored in your models library: {your_username}/{b_lora_trained_folder}"