Update app.py

app.py

@@ -160,6 +160,7 @@ def recursive_update(d, u):
 def start_training(
     lora_name,
     concept_sentence,
+    which_model,
     steps,
     lr,
     rank,
@@ -224,7 +225,12 @@ def start_training(
             config["config"]["process"][0]["sample"]["prompts"].append(sample_3)
     else:
         config["config"]["process"][0]["train"]["disable_sampling"] = True
-   
+        
+    if(which_model == "[schnell] (4 step fast model)"):
+        config["config"]["process"][0]["model"]["name_or_path"] = "black-forest-labs/FLUX.1-schnell"
+        config["config"]["process"][0]["model"]["assistant_lora_path"] = "ostris/FLUX.1-schnell-training-adapter"
+        config["config"]["process"][0]["sample"]["sample_steps"] = 4
+
     if(use_more_advanced_options):
         more_advanced_options_dict = yaml.safe_load(more_advanced_options)
         config["config"]["process"][0] = recursive_update(config["config"]["process"][0], more_advanced_options_dict)
@@ -291,11 +297,13 @@ def update_pricing(steps, oauth_token: Union[gr.OAuthToken, None]):
     else:
         return gr.update(visible=False), "", gr.update(visible=False), gr.update(visible=True)
 
+def swap_base_model(model):
+    return gr.update(visible=True) if model == "[dev] (high quality model, non-commercial license)" else gr.update(visible=False)    
+
 config_yaml = '''
 device: cuda:0
 model:
   is_flux: true
-  name_or_path: black-forest-labs/FLUX.1-dev
   quantize: true
 network:
   linear: 16 #it will overcome the 'rank' parameter
@@ -342,6 +350,7 @@ h3{margin-top: 0}
 .main_ui_logged_out{opacity: 0.3; pointer-events: none}
 .tabitem{border: 0px}
 .group_padding{padding: .55em}
+#space_model .wrap > label:last-child{opacity: 0.3; pointer-events:none}
 """
 with gr.Blocks(theme=theme, css=css) as demo:
     gr.Markdown(
@@ -352,18 +361,22 @@ with gr.Blocks(theme=theme, css=css) as demo:
         gr.LoginButton("Sign in with Hugging Face to train your LoRA on Spaces", visible=is_spaces)
     with gr.Tab("Train on Spaces" if is_spaces else "Train locally"):
         with gr.Column() as main_ui:
-            with gr.Row():
-                lora_name = gr.Textbox(
-                    label="The name of your LoRA",
-                    info="This has to be a unique name",
-                    placeholder="e.g.: Persian Miniature Painting style, Cat Toy",
-                )
-                concept_sentence = gr.Textbox(
-                    label="Trigger word/sentence",
-                    info="Trigger word or sentence to be used",
-                    placeholder="uncommon word like p3rs0n or trtcrd, or sentence like 'in the style of CNSTLL'",
-                    interactive=True,
-                )
+            with gr.Group():
+                with gr.Row():
+                    lora_name = gr.Textbox(
+                        label="The name of your LoRA",
+                        info="This has to be a unique name",
+                        placeholder="e.g.: Persian Miniature Painting style, Cat Toy",
+                    )
+                    concept_sentence = gr.Textbox(
+                        label="Trigger word/sentence",
+                        info="Trigger word or sentence to be used",
+                        placeholder="uncommon word like p3rs0n or trtcrd, or sentence like 'in the style of CNSTLL'",
+                        interactive=True,
+                    )
+                which_model = gr.Radio(["[schnell] (4 step fast model)", "[dev] (high quality model, non-commercial license - available when training locally)"], label="Which base model to train?", elem_id="space_model" if is_spaces else "local_model", value="[schnell] (4 step fast model)",)
+                model_warning = gr.Markdown("""> [dev] model license is non-commercial. By choosing to fine-tune [dev], you must agree with [its license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md) and make sure the LoRA you will train and the training process you would start does not violate it.
+                """, visible=False)
             with gr.Group(visible=True) as image_upload:
                 with gr.Row():
                     images = gr.File(
@@ -503,12 +516,18 @@ with gr.Blocks(theme=theme, css=css) as demo:
         inputs=[steps],
         outputs=[cost_preview, cost_preview_info, payment_update, start]
     )
-    
+
+    which_model.change(
+        fn=swap_base_model,
+        inputs=which_model,
+        outputs=model_warning
+    )
     start.click(fn=create_dataset, inputs=[images] + caption_list, outputs=dataset_folder).then(
         fn=start_training,
         inputs=[
             lora_name,
             concept_sentence,
+            which_model,
             steps,
             lr,
             rank,