Update app_training.py

app_training.py

@@ -13,6 +13,7 @@ from trainer import Trainer
 
 def create_training_demo(trainer: Trainer,
                          pipe: InferencePipeline | None = None) -> gr.Blocks:
+    hf_token = os.getenv('HF_TOKEN')
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
@@ -32,7 +33,8 @@ def create_training_demo(trainer: Trainer,
                                                 max_lines=1)
                     delete_existing_model = gr.Checkbox(
                         label='Delete existing model of the same name',
-                        value=False)
+                        value=False,
+                        visible=False)
                     validation_prompt = gr.Text(label='Validation Prompt')
                 with gr.Box():
                     gr.Markdown('Upload Settings')
@@ -63,31 +65,35 @@ def create_training_demo(trainer: Trainer,
                                              value='512',
                                              label='Resolution',
                                              visible=False)
-                num_training_steps = gr.Number(
-                    label='Number of Training Steps', value=300, precision=0)
-                learning_rate = gr.Number(label='Learning Rate',
-                                          value=0.000035)
-                gradient_accumulation = gr.Number(
-                    label='Number of Gradient Accumulation',
-                    value=1,
-                    precision=0)
-                seed = gr.Slider(label='Seed',
-                                 minimum=0,
-                                 maximum=100000,
-                                 step=1,
-                                 value=0)
-                fp16 = gr.Checkbox(label='FP16', value=True)
-                use_8bit_adam = gr.Checkbox(label='Use 8bit Adam', value=False)
-                checkpointing_steps = gr.Number(label='Checkpointing Steps',
-                                                value=1000,
-                                                precision=0)
-                validation_epochs = gr.Number(label='Validation Epochs',
-                                              value=100,
-                                              precision=0)
+                
+                token = gr.Text(label="Hugging Face Write Token", placeholder="", visible=True if hf_token else False)
+                with gr.Accordion("Advanced settings", open=False):
+                    num_training_steps = gr.Number(
+                        label='Number of Training Steps', value=300, precision=0)
+                    learning_rate = gr.Number(label='Learning Rate',
+                                              value=0.000035)
+                    gradient_accumulation = gr.Number(
+                        label='Number of Gradient Accumulation',
+                        value=1,
+                        precision=0)
+                    seed = gr.Slider(label='Seed',
+                                     minimum=0,
+                                     maximum=100000,
+                                     step=1,
+                                     randomize=True,
+                                     value=0)
+                    fp16 = gr.Checkbox(label='FP16', value=True)
+                    use_8bit_adam = gr.Checkbox(label='Use 8bit Adam', value=False)
+                    checkpointing_steps = gr.Number(label='Checkpointing Steps',
+                                                    value=1000,
+                                                    precision=0)
+                    validation_epochs = gr.Number(label='Validation Epochs',
+                                                  value=100,
+                                                  precision=0)
                 gr.Markdown('''
                     - The base model must be a model that is compatible with [diffusers](https://github.com/huggingface/diffusers) library.
                     - It takes a few minutes to download the base model first.
-                    - Expected time to train a model for 300 steps: 20 minutes with T4, 8 minutes with A10G, (4 minutes with A100)
+                    - Expected time to train a model for 300 steps: ~20 minutes with T4
                     - It takes a few minutes to upload your trained model.
                     - You may want to try a small number of steps first, like 1, to see if everything works fine in your environment.
                     - You can check the training status by pressing the "Open logs" button if you are running this on your Space.
@@ -128,6 +134,7 @@ def create_training_demo(trainer: Trainer,
                              delete_existing_repo,
                              upload_to,
                              remove_gpu_after_training,
+                             token
                          ],
                          outputs=output_message)
     return demo