update two buttons

app_training.py

@@ -40,6 +40,15 @@ def create_training_demo(trainer: Trainer,
                                                  value='512',
                                                  label='Resolution',
                                                  visible=False)
+                    with gr.Row():
+                        tuned_model = gr.Text(
+                            label='Path to tuned model',
+                            value='xxx/xxx,
+                            max_lines=1)
+                        resolution = gr.Dropdown(choices=['512', '768'],
+                                                 value='512',
+                                                 label='Resolution',
+                                                 visible=False)
 
                     input_token = gr.Text(label='Hugging Face Write Token',
                                           placeholder='',
@@ -153,7 +162,7 @@ def create_training_demo(trainer: Trainer,
                 gradient_accumulation, seed, fp16, use_8bit_adam,
                 checkpointing_steps, validation_epochs, upload_to_hub,
                 use_private_repo, delete_existing_repo, upload_to,
-                remove_gpu_after_training, input_token, blend_word_1, blend_word_2, eq_params_1, eq_params_2
+                remove_gpu_after_training, input_token, blend_word_1, blend_word_2, eq_params_1, eq_params_2, tuned_model
             ],
             outputs=output_message)
     return demo

trainer.py

@@ -207,6 +207,7 @@ class Trainer:
         blend_word_2: str,
         eq_params_1: str,
         eq_params_2: str,
+        tuned_model: str = None,
     ) -> str:
         # if SPACE_ID == ORIGINAL_SPACE_ID:
         #     raise gr.Error(
@@ -239,7 +240,7 @@ class Trainer:
                 self.hf_token if self.hf_token else input_token)
 
         config = OmegaConf.load('Video-P2P/configs/man-skiing.yaml')
-        config.pretrained_model_path = self.download_base_model(base_model)
+        config.pretrained_model_path = self.download_base_model(tuned_model)
         config.output_dir = output_dir.as_posix()
         config.train_data.video_path = training_video.name  # type: ignore
         config.train_data.prompt = training_prompt