Update app.py

app.py

@@ -82,31 +82,18 @@ def download_file(url, directory=None):
     
     return filepath
             
-def update_selection(evt: gr.SelectData, selected_indices, width, height):
+def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, height):
     selected_index = evt.index
     selected_indices = selected_indices or []
     if selected_index in selected_indices:
-        # LoRA is already selected, remove it
         selected_indices.remove(selected_index)
     else:
         if len(selected_indices) < 2:
             selected_indices.append(selected_index)
         else:
             gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
-            return (
-                gr.update(),
-                gr.update(),
-                gr.update(),
-                gr.update(),
-                gr.update(),
-                gr.update(),
-                gr.update(),
-                gr.update(),
-                gr.update(),
-                gr.update(),
-            )
+            return gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), width, height, gr.update(), gr.update()
 
-    # Initialize outputs
     selected_info_1 = "Select a LoRA 1"
     selected_info_2 = "Select a LoRA 2"
     lora_scale_1 = 1.15
@@ -114,39 +101,25 @@ def update_selection(evt: gr.SelectData, selected_indices, width, height):
     lora_image_1 = None
     lora_image_2 = None
     if len(selected_indices) >= 1:
-        lora1 = loras[selected_indices[0]]
+        lora1 = loras_state[selected_indices[0]]
         selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
         lora_image_1 = lora1['image']
     if len(selected_indices) >= 2:
-        lora2 = loras[selected_indices[1]]
+        lora2 = loras_state[selected_indices[1]]
         selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
         lora_image_2 = lora2['image']
 
-    # Update prompt placeholder based on last selected LoRA
     if selected_indices:
-        last_selected_lora = loras[selected_indices[-1]]
+        last_selected_lora = loras_state[selected_indices[-1]]
         new_placeholder = f"Type a prompt for {last_selected_lora['title']}"
     else:
         new_placeholder = "Type a prompt after selecting a LoRA"
 
-    return (
-        gr.update(placeholder=new_placeholder),
-        selected_info_1,
-        selected_info_2,
-        selected_indices,
-        lora_scale_1,
-        lora_scale_2,
-        width,
-        height,
-        lora_image_1,
-        lora_image_2,
-    )
+    return gr.update(placeholder=new_placeholder), selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2
 
-def remove_lora_1(selected_indices):
-    selected_indices = selected_indices or []
+def remove_lora_1(selected_indices, loras_state):
     if len(selected_indices) >= 1:
         selected_indices.pop(0)
-    # Update selected_info_1 and selected_info_2
     selected_info_1 = "Select a LoRA 1"
     selected_info_2 = "Select a LoRA 2"
     lora_scale_1 = 1.15
@@ -154,20 +127,18 @@ def remove_lora_1(selected_indices):
     lora_image_1 = None
     lora_image_2 = None
     if len(selected_indices) >= 1:
-        lora1 = loras[selected_indices[0]]
+        lora1 = loras_state[selected_indices[0]]
         selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
         lora_image_1 = lora1['image']
     if len(selected_indices) >= 2:
-        lora2 = loras[selected_indices[1]]
+        lora2 = loras_state[selected_indices[1]]
         selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
         lora_image_2 = lora2['image']
     return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
 
-def remove_lora_2(selected_indices):
-    selected_indices = selected_indices or []
+def remove_lora_2(selected_indices, loras_state):
     if len(selected_indices) >= 2:
         selected_indices.pop(1)
-    # Update selected_info_1 and selected_info_2
     selected_info_1 = "Select a LoRA 1"
     selected_info_2 = "Select a LoRA 2"
     lora_scale_1 = 1.15
@@ -175,21 +146,21 @@ def remove_lora_2(selected_indices):
     lora_image_1 = None
     lora_image_2 = None
     if len(selected_indices) >= 1:
-        lora1 = loras[selected_indices[0]]
+        lora1 = loras_state[selected_indices[0]]
         selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
         lora_image_1 = lora1['image']
     if len(selected_indices) >= 2:
-        lora2 = loras[selected_indices[1]]
+        lora2 = loras_state[selected_indices[1]]
         selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
         lora_image_2 = lora2['image']
     return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
 
-def randomize_loras(selected_indices):
-    if len(loras) < 2:
+def randomize_loras(selected_indices, loras_state):
+    if len(loras_state) < 2:
         raise gr.Error("Not enough LoRAs to randomize.")
-    selected_indices = random.sample(range(len(loras)), 2)
-    lora1 = loras[selected_indices[0]]
-    lora2 = loras[selected_indices[1]]
+    selected_indices = random.sample(range(len(loras_state)), 2)
+    lora1 = loras_state[selected_indices[0]]
+    lora2 = loras_state[selected_indices[1]]
     selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
     selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
     lora_scale_1 = 1.15
@@ -199,6 +170,106 @@ def randomize_loras(selected_indices):
     random_prompt = random.choice(prompt_values)
     return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, random_prompt
 
+def add_custom_lora(custom_lora, selected_indices, current_loras):
+    if custom_lora:
+        try:
+            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
+            print(f"Loaded custom LoRA: {repo}")
+            existing_item_index = next((index for (index, item) in enumerate(current_loras) if item['repo'] == repo), None)
+            if existing_item_index is None:
+                if repo.endswith(".safetensors") and repo.startswith("http"):
+                    repo = download_file(repo)
+                new_item = {
+                    "image": image if image else "/home/user/app/custom.png",
+                    "title": title,
+                    "repo": repo,
+                    "weights": path,
+                    "trigger_word": trigger_word
+                }
+                print(f"New LoRA: {new_item}")
+                existing_item_index = len(current_loras)
+                current_loras.append(new_item)
+            
+            # Update gallery
+            gallery_items = [(item["image"], item["title"]) for item in current_loras]
+            # Update selected_indices if there's room
+            if len(selected_indices) < 2:
+                selected_indices.append(existing_item_index)
+            else:
+                gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
+
+            # Update selected_info and images
+            selected_info_1 = "Select a LoRA 1"
+            selected_info_2 = "Select a LoRA 2"
+            lora_scale_1 = 1.15
+            lora_scale_2 = 1.15
+            lora_image_1 = None
+            lora_image_2 = None
+            if len(selected_indices) >= 1:
+                lora1 = current_loras[selected_indices[0]]
+                selected_info_1 = f"### LoRA 1 Selected: {lora1['title']} ✨"
+                lora_image_1 = lora1['image'] if lora1['image'] else None
+            if len(selected_indices) >= 2:
+                lora2 = current_loras[selected_indices[1]]
+                selected_info_2 = f"### LoRA 2 Selected: {lora2['title']} ✨"
+                lora_image_2 = lora2['image'] if lora2['image'] else None
+            print("Finished adding custom LoRA")
+            return (
+                current_loras,
+                gr.update(value=gallery_items),
+                selected_info_1, 
+                selected_info_2,
+                selected_indices,
+                lora_scale_1,
+                lora_scale_2,
+                lora_image_1,
+                lora_image_2
+            )
+        except Exception as e:
+            print(e)
+            gr.Warning(str(e))
+            return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
+    else:
+        return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
+
+def remove_custom_lora(selected_indices, current_loras):
+    if current_loras:
+        custom_lora_repo = current_loras[-1]['repo']
+        # Remove from loras list
+        current_loras = current_loras[:-1]
+        # Remove from selected_indices if selected
+        custom_lora_index = len(current_loras)
+        if custom_lora_index in selected_indices:
+            selected_indices.remove(custom_lora_index)
+    # Update gallery
+    gallery_items = [(item["image"], item["title"]) for item in current_loras]
+    # Update selected_info and images
+    selected_info_1 = "Select a LoRA 1"
+    selected_info_2 = "Select a LoRA 2"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_image_1 = None
+    lora_image_2 = None
+    if len(selected_indices) >= 1:
+        lora1 = current_loras[selected_indices[0]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
+        lora_image_1 = lora1['image']
+    if len(selected_indices) >= 2:
+        lora2 = current_loras[selected_indices[1]]
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
+        lora_image_2 = lora2['image']
+    return (
+        current_loras,
+        gr.update(value=gallery_items),
+        selected_info_1,
+        selected_info_2,
+        selected_indices,
+        lora_scale_1,
+        lora_scale_2,
+        lora_image_1,
+        lora_image_2
+    )
+
 @spaces.GPU(duration=70)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
     print("Generating image...")
@@ -206,8 +277,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
     generator = torch.Generator(device="cuda").manual_seed(seed)
     with calculateDuration("Generating image"):
         # Generate image
-        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-            prompt=prompt_mash,
+        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(prompt=prompt_mash,
             num_inference_steps=steps,
             guidance_scale=cfg_scale,
             width=width,
@@ -238,11 +308,12 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     ).images[0]
     return final_image
 
-def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, progress=gr.Progress(track_tqdm=True)):
+@spaces.GPU(duration=70)
+def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
     if not selected_indices:
         raise gr.Error("You must select at least one LoRA before proceeding.")
 
-    selected_loras = [loras[idx] for idx in selected_indices]
+    selected_loras = [loras_state[idx] for idx in selected_indices]
 
     # Build the prompt with trigger words
     prepends = []
@@ -361,108 +432,6 @@ def check_custom_model(link):
         # Assume it's a Hugging Face model path
         return get_huggingface_safetensors(link)
 
-def add_custom_lora(custom_lora, selected_indices):
-    global loras
-    if custom_lora:
-        try:
-            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
-            print(f"Loaded custom LoRA: {repo}")
-            existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
-            if existing_item_index is None:
-                if(repo.endswith(".safetensors") and repo.startswith("http")):
-                    repo = download_file(repo)
-                new_item = {
-                    "image": image if image else "/home/user/app/custom.png",
-                    "title": title,
-                    "repo": repo,
-                    "weights": path,
-                    "trigger_word": trigger_word
-                }
-                print(f"New LoRA: {new_item}")
-                existing_item_index = len(loras)
-                loras.append(new_item)
-            
-            # Update gallery
-            gallery_items = [(item["image"], item["title"]) for item in loras]
-            # Update selected_indices if there's room
-            if len(selected_indices) < 2:
-                selected_indices.append(existing_item_index)
-            else:
-                gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
-
-            # Update selected_info and images
-            selected_info_1 = "Select a LoRA 1"
-            selected_info_2 = "Select a LoRA 2"
-            lora_scale_1 = 1.15
-            lora_scale_2 = 1.15
-            lora_image_1 = None
-            lora_image_2 = None
-            if len(selected_indices) >= 1:
-                lora1 = loras[selected_indices[0]]
-                selected_info_1 = f"### LoRA 1 Selected: {lora1['title']} ✨"
-                lora_image_1 = lora1['image'] if lora1['image'] else None
-            if len(selected_indices) >= 2:
-                lora2 = loras[selected_indices[1]]
-                selected_info_2 = f"### LoRA 2 Selected: {lora2['title']} ✨"
-                lora_image_2 = lora2['image'] if lora2['image'] else None
-            print("Chegou no final")
-            return (
-                gr.update(value=gallery_items),
-                selected_info_1, 
-                selected_info_2,
-                selected_indices,
-                lora_scale_1,
-                lora_scale_2,
-                lora_image_1,
-                lora_image_2
-            )
-        except Exception as e:
-            print(e)
-            gr.Warning(str(e))
-            return gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
-    else:
-        return gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
-
-def remove_custom_lora(selected_indices):
-    global loras
-    if loras:
-        custom_lora_repo = loras[-1]['repo']
-        # Remove from loras list
-        loras = loras[:-1]
-        # Remove from selected_indices if selected
-        custom_lora_index = len(loras)
-        if custom_lora_index in selected_indices:
-            selected_indices.remove(custom_lora_index)
-    # Update gallery
-    gallery_items = [(item["image"], item["title"]) for item in loras]
-    # Update selected_info and images
-    selected_info_1 = "Select a LoRA 1"
-    selected_info_2 = "Select a LoRA 2"
-    lora_scale_1 = 1.15
-    lora_scale_2 = 1.15
-    lora_image_1 = None
-    lora_image_2 = None
-    if len(selected_indices) >= 1:
-        lora1 = loras[selected_indices[0]]
-        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
-        lora_image_1 = lora1['image']
-    if len(selected_indices) >= 2:
-        lora2 = loras[selected_indices[1]]
-        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
-        lora_image_2 = lora2['image']
-    return (
-        gr.update(value=gallery_items),
-        selected_info_1,
-        selected_info_2,
-        selected_indices,
-        lora_scale_1,
-        lora_scale_2,
-        lora_image_1,
-        lora_image_2
-    )
-
-
-
 css = '''
 #gen_btn{height: 100%}
 #title{text-align: center}
@@ -484,6 +453,7 @@ css = '''
 #custom_lora_btn{margin-top: auto;margin-bottom: 11px}
 #random_btn{font-size: 300%}
 '''
+
 with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
     title = gr.HTML(
         """<h1><img src="https://i.imgur.com/wMh2Oek.png" alt="LoRA"> LoRA Lab [beta]</h1><br><span style="
@@ -493,6 +463,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
 ">Mix and match any FLUX[dev] LoRAs</span>""",
         elem_id="title",
     )
+    loras_state = gr.State(loras)
     selected_indices = gr.State([])
     with gr.Row():
         with gr.Column(scale=3):
@@ -560,40 +531,39 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
 
     gallery.select(
         update_selection,
-        inputs=[selected_indices, width, height],
-        outputs=[prompt, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2]
-    )
+        inputs=[selected_indices, loras_state, width, height],
+        outputs=[prompt, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2])
     remove_button_1.click(
         remove_lora_1,
-        inputs=[selected_indices],
+        inputs=[selected_indices, loras_state],
         outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
     remove_button_2.click(
         remove_lora_2,
-        inputs=[selected_indices],
+        inputs=[selected_indices, loras_state],
         outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
     randomize_button.click(
         randomize_loras,
-        inputs=[selected_indices],
+        inputs=[selected_indices, loras_state],
         outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, prompt]
     )
     add_custom_lora_button.click(
         add_custom_lora,
-        inputs=[custom_lora, selected_indices],
-        outputs=[gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
+        inputs=[custom_lora, selected_indices, loras_state],
+        outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
     remove_custom_lora_button.click(
         remove_custom_lora,
-        inputs=[selected_indices],
-        outputs=[gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
+        inputs=[selected_indices, loras_state],
+        outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height],
+        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
         outputs=[result, seed, progress_bar]
     )
 
 app.queue()
-app.launch()
+app.launch()