ready to release

app.py

@@ -59,12 +59,12 @@ pipe = InferencePipeline(HF_TOKEN)
 trainer = Trainer(HF_TOKEN)
 
 with gr.Blocks(css='style.css') as demo:
-    # if SPACE_ID == ORIGINAL_SPACE_ID:
-    #     show_warning(SHARED_UI_WARNING)
-    # elif not torch.cuda.is_available():
-    #     show_warning(CUDA_NOT_AVAILABLE_WARNING)
-    # elif (not 'T4' in GPU_DATA):
-    #     show_warning(INVALID_GPU_WARNING)
+    if SPACE_ID == ORIGINAL_SPACE_ID:
+        show_warning(SHARED_UI_WARNING)
+    elif not torch.cuda.is_available():
+        show_warning(CUDA_NOT_AVAILABLE_WARNING)
+    elif (not 'T4' in GPU_DATA):
+        show_warning(INVALID_GPU_WARNING)
 
     gr.Markdown(TITLE)
     with gr.Tabs():

app_training.py

@@ -43,7 +43,7 @@ def create_training_demo(trainer: Trainer,
                     with gr.Row():
                         tuned_model = gr.Text(
                             label='Path to tuned model',
-                            value='xxx/xxx',
+                            value='xxx/ski-lego',
                             max_lines=1)
                         resolution = gr.Dropdown(choices=['512', '768'],
                                                  value='512',
@@ -60,6 +60,8 @@ def create_training_demo(trainer: Trainer,
                             precision=0)
                         learning_rate = gr.Number(label='Learning Rate',
                                                   value=0.000035)
+                        cross_replace = gr.Number(label='Cross attention replace ratio',
+                                                  value=0.2)
                         gradient_accumulation = gr.Number(
                             label='Number of Gradient Accumulation',
                             value=1,
@@ -78,7 +80,7 @@ def create_training_demo(trainer: Trainer,
                             value=1000,
                             precision=0)
                         validation_epochs = gr.Number(
-                            label='Validation Epochs', value=100, precision=0)
+                            label='Validation Epochs', value=300, precision=0)
                     gr.Markdown('''
                         - The base model must be a Stable Diffusion model compatible with [diffusers](https://github.com/huggingface/diffusers) library.
                         - Expected time to complete: ~20 minutes with T4.
@@ -89,8 +91,8 @@ def create_training_demo(trainer: Trainer,
         with gr.Row():
             with gr.Column():
                 gr.Markdown('Output Model')
-                output_model_name = gr.Text(label='Name of your model',
-                                            placeholder='The skiing man',
+                output_model_name = gr.Text(label='Path to save your tuned model',
+                                            placeholder='ski-lego',
                                             max_lines=1)
                 validation_prompt = gr.Text(
                     label='Validation Prompt',
@@ -111,7 +113,7 @@ def create_training_demo(trainer: Trainer,
                 eq_params_2 = gr.Text(
                     label='reweight_value',
                     placeholder=
-                    '8')
+                    '4')
             with gr.Column():
                 gr.Markdown('Upload Settings')
                 with gr.Row():
@@ -162,7 +164,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, tuned_model
+                remove_gpu_after_training, input_token, blend_word_1, blend_word_2, eq_params_1, eq_params_2, tuned_model, cross_replace
             ],
             outputs=output_message)
     return demo

trainer.py

@@ -217,6 +217,7 @@ class Trainer:
         eq_params_1: str,
         eq_params_2: str,
         tuned_model: str,
+        cross_replace: float,
     ) -> str:
         # if SPACE_ID == ORIGINAL_SPACE_ID:
         #     raise gr.Error(
@@ -280,6 +281,7 @@ class Trainer:
             config.is_word_swap = True
         else:
             config.is_word_swap = False
+        config.cross_replace_steps = cross_replace
 
         config_path = output_dir / 'config.yaml'
         with open(config_path, 'w') as f: