Update app.py

app.py

@@ -37,16 +37,16 @@ def swap_hardware(hf_token, hardware="cpu-basic"):
     requests.post(hardware_url, json = body, headers=headers)
 
 
-global sleep_time
 
-def swap_sleep_time(hf_token,sleep_time):
+
+def swap_sleep_time(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):
+def get_sleep_time():
     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)
@@ -54,10 +54,15 @@ def get_sleep_time(hf_token):
         gcTimeout = response.json()['runtime']['gcTimeout']
     except:
         gcTimeout = None
-    sleep_time = gcTimeout
+    
     return gcTimeout
 
-sleep_time = get_sleep_time(hf_token)
+def check_sleep_time():
+    sleep_time = get_sleep_time()
+    if sleep_time is None :
+        return sleep_time, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+    else :
+        return sleep_time, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
     
 def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instance_prompt, class_prompt, max_train_steps, checkpoint_steps):
     
@@ -159,8 +164,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_sleep_time(hf_token, 300)
+    #swap_hardware(hardware="cpu-basic")
+    swap_sleep_time(300)
     
     return f"Done, your trained model has been stored in your models library: {your_username}/{b_lora_trained_folder}"
 
@@ -264,40 +269,29 @@ with gr.Blocks(css=css) as demo:
             image = gr.Image(label="Image Reference", sources=["upload"], type="filepath")
 
             with gr.Column():
-                current_sleep_time = gr.Number(visible=False)
-                print(f"loaded demo: {sleep_time}")
-                if sleep_time is not None and sleep_time > 1:
-                    gr.HTML(f'''
-                    <div class="gr-prose">
-                    <h2>in order to train properly your B-LoRa, you need to set space sleep time to "Don't sleep"</h2>
-                    <p class="actions">
-                        <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings">⚡️ &nbsp; Set recommended Sleep Time</a>
-                    </p>
-                    </div>
-                    ''')
-                    training_type = gr.Radio(label="Training type", choices=["style", "concept"], value="style", interactive=False)
-                    b_lora_name = gr.Textbox(label="Name your B-LoRa model", placeholder="b_lora_trained_folder", interactive=False)
+                
                     with gr.Row():
-                        instance_prompt = gr.Textbox(label="Create instance prompt", placeholder="A [v42] <class_prompt>", interactive=False)
-                        class_prompt = gr.Textbox(label="Specify class prompt", placeholder="style | person | dog ", interactive=False)
-                    training_steps = gr.Number(label="Training steps", value=1000, interactive=False)
-                    checkpoint_step = gr.Number(label="checkpoint step", visible=False, value=500)
-                else:    
-                    training_type = gr.Radio(label="Training type", choices=["style", "concept"], value="style", interactive=True)
-                    b_lora_name = gr.Textbox(label="Name your B-LoRa model", placeholder="b_lora_trained_folder", interactive=True)
+                        current_sleep_time = gr.Textbox(label="current space sleep time")
+                        check_sleep_time_btn = gr.Button("check", scale=1)
+                    training_type = gr.Radio(label="Training type", choices=["style", "concept"], value="style", visible=False)
+                    b_lora_name = gr.Textbox(label="Name your B-LoRa model", placeholder="b_lora_trained_folder", visible=False)
                     with gr.Row():
-                        instance_prompt = gr.Textbox(label="Create instance prompt", placeholder="A [v42] <class_prompt>", interactive=True)
-                        class_prompt = gr.Textbox(label="Specify class prompt", placeholder="style | person | dog ", interactive=True)
-                    training_steps = gr.Number(label="Training steps", value=1000, interactive=False)
+                        instance_prompt = gr.Textbox(label="Create instance prompt", placeholder="A [v42] <class_prompt>", visible=False)
+                        class_prompt = gr.Textbox(label="Specify class prompt", placeholder="style | person | dog ", visible=False)
+                    training_steps = gr.Number(label="Training steps", value=1000, interactive=False, visible=False)
                     checkpoint_step = gr.Number(label="checkpoint step", visible=False, value=500)
 
                     
-        if sleep_time is not None and sleep_time > 1:
-            train_btn = gr.Button("Train B-LoRa", visible=False)
-        else:
-            train_btn = gr.Button("Train B-LoRa", visible=True)
+        
+        train_btn = gr.Button("Train B-LoRa", visible=False)
+    
         status = gr.Textbox(label="status")
 
+    check_sleep_time_btn.click(
+        fn = check_sleep_time,
+        inputs = None, 
+        outputs = [current_sleep_time, b_lora_name, instance_prompt, class_prompt, training_type, training_steps, train_btn]
+    )
     
     training_type.change(
                         fn = change_training_setup,
@@ -310,5 +304,5 @@ with gr.Blocks(css=css) as demo:
         inputs = [image, b_lora_name, instance_prompt, class_prompt, training_type, training_steps],
         outputs = [status]
     )
-    demo.load(fn=get_sleep_time, inputs=None, outputs=[current_sleep_time])
+    
 demo.launch(debug=True)