Update app.py

app.py

@@ -32,7 +32,11 @@ def start_train():
     # Handles CUDA OOM errors.
     os.system(f"export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True")
     os.system("echo 'Okay, trying training.'")
+    status_box.value = "Status: Training"
+    API.add_space_variable(repo_id=experiment_repo, key="STATUS", value="TRAINING")
     os.system(f"cd pytorch-image-models; ./train.sh 4 --dataset hfds/datacomp/imagenet-1k-random-10.0-frac-1over4 --log-wandb --wandb-project ImageNetTraining10.0-frac-1over4 --experiment ImageNetTraining10.0-frac-1over4 --model seresnet34 --sched cosine --epochs 150 --warmup-epochs 5 --lr 0.4 --reprob 0.5 --remode pixel --batch-size 256 --amp -j 4")
+    API.add_space_variable(repo_id=experiment_repo, key="STATUS", value="IDLE")
+    status_box.value = "Status: Idle"
     os.system("echo 'Done'.")
     os.system("ls")
     # Upload output to repository
@@ -44,8 +48,15 @@ def run():
     with gr.Blocks() as app:
         gr.Markdown(f"Randomization: 10.0")
         gr.Markdown(f"Subset: frac-1over4")
-        start = gr.Button("Start")
-        start.click(start_train)
+        status_box = gr.Textbox()
+        space_variables = API.get_space_variables(repo_id=experiment_repo)
+        if 'STATUS' not in space_variables or space_variables['STATUS'].value != 'TRAINING':
+            API.add_space_variable(repo_id=experiment_repo, key="STATUS", value="IDLE")
+            status_box.value = "Status: Idle"
+            start = gr.Button("Start")
+            start.click(start_train, inputs=status_box)
+        else:
+            status_box.value = "Status: Training"
     app.launch(server_name="0.0.0.0", server_port=7860)
     
 if __name__ == '__main__':