Update app.py

app.py

@@ -39,6 +39,10 @@ MAX_IMAGES = 150
 
 
 def load_captioning(uploaded_images, concept_sentence):
+    
+    uploaded_images = [file for file in uploaded_images if not file.endswith('.txt')]
+    txt_files = [file for file in uploaded_images if file.endswith('.txt')]
+    
     gr.Info("Images uploaded!")
     updates = []
     if len(uploaded_images) <= 1:
@@ -54,17 +58,23 @@ def load_captioning(uploaded_images, concept_sentence):
     for i in range(1, MAX_IMAGES + 1):
         # Determine if the current row and image should be visible
         visible = i <= len(uploaded_images)
-
+        
         # Update visibility of the captioning row
         updates.append(gr.update(visible=visible))
 
         # Update for image component - display image if available, otherwise hide
         image_value = uploaded_images[i - 1] if visible else None
-
         updates.append(gr.update(value=image_value, visible=visible))
 
+        base_name = image_value.rsplit('.', 1)[0]
+        
+        corresponding_txt = base_name + '.txt'
+        corresponding_caption = False
+        if corresponding_txt in txt_files:
+            with open(corresponding_txt, 'r') as file:
+                corresponding_caption = file.read()
         # Update value of captioning area
-        text_value = "[trigger]" if visible and concept_sentence else None
+        text_value = corresponding_caption if corresponding_caption else "[trigger]" if visible and concept_sentence else None
         updates.append(gr.update(value=text_value, visible=visible))
 
     # Update for the sample caption area
@@ -145,6 +155,8 @@ def start_training(
     sample_1,
     sample_2,
     sample_3,
+    use_more_advanced_options,
+    more_advanced_options,
     profile: Union[gr.OAuthProfile, None],
     oauth_token: Union[gr.OAuthToken, None],
 ):
@@ -197,6 +209,10 @@ def start_training(
             config["config"]["process"][0]["sample"]["prompts"].append(sample_3)
     else:
         config["config"]["process"][0]["train"]["disable_sampling"] = True
+   
+    if(use_more_advanced_options):
+        config["config"]["process"] = more_advanced_options
+    
     # Save the updated config
     # generate a random name for the config
     random_config_name = str(uuid.uuid4())
@@ -232,20 +248,6 @@ def start_training(
 
     return f"Training completed successfully. Model saved as {slugged_lora_name}"
 
-
-theme = gr.themes.Monochrome(
-    text_size=gr.themes.Size(lg="18px", md="15px", sm="13px", xl="22px", xs="12px", xxl="24px", xxs="9px"),
-    font=[gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui", "sans-serif"],
-)
-css = """
-h1{font-size: 2em}
-h3{margin-top: 0}
-#component-1{text-align:center}
-.main_ui_logged_out{opacity: 0.3; pointer-events: none}
-.tabitem{border: 0px}
-.group_padding{padding: .55em}
-"""
-
 def swap_visibilty(profile: Union[gr.OAuthProfile, None]):
     if is_spaces:
         if profile is None:
@@ -272,6 +274,64 @@ def update_pricing(steps, oauth_token: Union[gr.OAuthToken, None]):
     else:
         return gr.update(visible=False), "", gr.update(visible=False), gr.update(visible=True)
 
+config_yaml = {
+    "device": "cuda:0",
+    "network": {
+        "type": "lora",
+        "linear": 16,
+        "linear_alpha": 16
+    },
+    "save": {
+        "dtype": "float16",
+        "save_every": 10000,
+        "max_step_saves_to_keep": 4,
+        "push_to_hub": True,
+        "hf_private": True
+    },
+    "train": {
+        "batch_size": 1,
+        "gradient_accumulation_steps": 1,
+        "train_unet": True,
+        "train_text_encoder": False,
+        "gradient_checkpointing": True,
+        "noise_scheduler": "flowmatch",
+        "optimizer": "adamw8bit",
+        "ema_config": {
+            "use_ema": True,
+            "ema_decay": 0.99
+        },
+        "dtype": "bf16"
+    },
+    "model": {
+        "name_or_path": "black-forest-labs/FLUX.1-dev",
+        "is_flux": True,
+        "quantize": True
+    },
+    "sample": {
+        "sampler": "flowmatch",
+        "sample_every": 1000,
+        "width": 1024,
+        "height": 1024,
+        "neg": "",
+        "seed": 42,
+        "walk_seed": True,
+        "guidance_scale": 3.5,
+        "sample_steps": 28
+    }
+}
+
+theme = gr.themes.Monochrome(
+    text_size=gr.themes.Size(lg="18px", md="15px", sm="13px", xl="22px", xs="12px", xxl="24px", xxs="9px"),
+    font=[gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui", "sans-serif"],
+)
+css = """
+h1{font-size: 2em}
+h3{margin-top: 0}
+#component-1{text-align:center}
+.main_ui_logged_out{opacity: 0.3; pointer-events: none}
+.tabitem{border: 0px}
+.group_padding{padding: .55em}
+"""
 with gr.Blocks(theme=theme, css=css) as demo:
     gr.Markdown(
         """# LoRA Ease for FLUX 🧞‍♂️
@@ -295,8 +355,9 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 )
             with gr.Group(visible=True) as image_upload:
                 with gr.Row():
+                    gr.Markdown("Upload your images to caption them in the UI (if you already have a dataset with .txt captions, upload them together)")
                     images = gr.File(
-                        file_types=["image"],
+                        file_types=["image", ".txt"],
                         label="Upload your images",
                         file_count="multiple",
                         interactive=True,
@@ -339,6 +400,11 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 steps = gr.Number(label="Steps", value=1000, minimum=1, maximum=10000, step=1)
                 lr = gr.Number(label="Learning Rate", value=4e-4, minimum=1e-6, maximum=1e-3, step=1e-6)
                 rank = gr.Number(label="LoRA Rank", value=16, minimum=4, maximum=128, step=4)
+                with gr.Accordion("Even more advanced options", open=False):
+                    if(is_spaces):
+                        gr.Markdown("Attention: changing this parameters may make your training fail or go out-of-memory if training on Spaces. Only change settings here it if you know what you are doing. Beware that training is done in an L4 GPU with 24GB of RAM")
+                    use_more_advanced_options = gr.Checkbox(label="Use more advanced options", value=False)
+                    more_advanced_options = gr.Code(config_yaml, language="yaml")
 
             with gr.Accordion("Sample prompts (optional)", visible=False) as sample:
                 gr.Markdown(
@@ -424,6 +490,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
             sample_1,
             sample_2,
             sample_3,
+            use_more_advanced_options,
+            more_advanced_options
         ],
         outputs=progress_area,
     )