Upload 2 files

app.py

@@ -18,7 +18,7 @@ import pandas as pd
 import numpy as np
 from pathlib import Path
 
-from env import models, models_dev, models_schnell, models_fill, models_canny, models_depth, models_edit, num_loras, num_cns, HF_TOKEN, single_file_base_models
+from env import models, models_dev, models_schnell, models_fill, models_canny, models_depth, models_edit, num_loras, num_cns, MAX_LORA, HF_TOKEN, single_file_base_models
 from mod import (clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists, get_model_trigger,
                  description_ui, compose_lora_json, is_valid_lora, fuse_loras, turbo_loras, save_image, preprocess_i2i_image,
                  get_trigger_word, enhance_prompt, set_control_union_image, get_canny_image, get_depth_image,
@@ -236,93 +236,52 @@ def download_file(url, directory=None):
     
     return filepath
 
+def get_lora_selected_info(selected_indices, current_loras):
+    selected_info = [f"Select a LoRA {int(i+1)}" for i in range(MAX_LORA)]
+    lora_scale = [1.15] * MAX_LORA
+    lora_image = [None] * MAX_LORA
+    for i in range(MAX_LORA):
+        if len(selected_indices) >= i + 1:
+            lora = current_loras[selected_indices[i]]
+            selected_info[i] = f"### LoRA {i+1} Selected: [{lora['title']}]({lora['repo']}) ✨"
+            lora_image[i] = lora['image'] if lora['image'] else None
+    return selected_info, lora_scale, lora_image
+
 def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, height):
+    updates = [gr.update()] * MAX_LORA
     selected_index = evt.index
     selected_indices = selected_indices or []
     if selected_index in selected_indices:
         selected_indices.remove(selected_index)
     else:
-        if len(selected_indices) < 2:
+        if len(selected_indices) < MAX_LORA:
             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(), selected_indices, gr.update(), gr.update(), width, height, gr.update(), gr.update()
-
-    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_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_state[selected_indices[1]]
-        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
-        lora_image_2 = lora2['image']
+            gr.Warning(f"You can select up to {MAX_LORA} LoRAs, remove one to select a new one.")
+            return gr.update(), *updates, selected_indices, *updates, width, height, *updates
 
+    selected_info, lora_scale, lora_image = get_lora_selected_info(selected_indices, loras_state)
     if selected_indices:
         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"
 
-    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, loras_state):
-    if len(selected_indices) >= 1:
-        selected_indices.pop(0)
-    selected_info_1 = "Select LoRA 1"
-    selected_info_2 = "Select 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_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_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, loras_state):
-    if len(selected_indices) >= 2:
-        selected_indices.pop(1)
-    selected_info_1 = "Select LoRA 1"
-    selected_info_2 = "Select 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_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_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
+    return gr.update(placeholder=new_placeholder), *selected_info, selected_indices, *lora_scale, width, height, *lora_image
+
+def remove_lora(selected_indices, loras_state, n: int):
+    if len(selected_indices) >= n + 1:
+        selected_indices.pop(n)
+    selected_info, lora_scale, lora_image = get_lora_selected_info(selected_indices, loras_state)
+    return *selected_info, selected_indices, *lora_scale, *lora_image
 
 def randomize_loras(selected_indices, loras_state):
-    if len(loras_state) < 2:
+    if len(loras_state) < MAX_LORA:
         raise gr.Error("Not enough LoRAs to randomize.")
-    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
-    lora_scale_2 = 1.15
-    lora_image_1 = lora1['image']
-    lora_image_2 = lora2['image']
+    selected_indices = random.sample(range(len(loras_state)), MAX_LORA)
     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
+    selected_info, lora_scale, lora_image = get_lora_selected_info(selected_indices, loras_state)
+    return *selected_info, selected_indices, *lora_scale, *lora_image, random_prompt
 
 def download_loras_images(loras_json_orig: list[dict]):
     api = HfApi(token=HF_TOKEN)
@@ -348,6 +307,7 @@ def download_loras_images(loras_json_orig: list[dict]):
     return loras_json
 
 def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
+    updates = [gr.update()] * MAX_LORA
     if custom_lora:
         try:
             title, repo, path, trigger_word, image = check_custom_model(custom_lora)
@@ -377,44 +337,28 @@ def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
             # 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:
+            if len(selected_indices) < MAX_LORA:
                 selected_indices.append(existing_item_index)
             else:
-                gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
+                gr.Warning(f"You can select up to {MAX_LORA} 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
+            selected_info, lora_scale, lora_image = get_lora_selected_info(selected_indices, current_loras)
             print("Finished adding custom LoRA")
             return (
                 current_loras,
                 gr.update(value=gallery_items),
-                selected_info_1, 
-                selected_info_2,
+                *selected_info,
                 selected_indices,
-                lora_scale_1,
-                lora_scale_2,
-                lora_image_1,
-                lora_image_2
+                *lora_scale,
+                *lora_image
             )
         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()
+            return current_loras, gr.update(), *updates, selected_indices, *updates, *updates
     else:
-        return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update()
+        return current_loras, gr.update(), *updates, selected_indices, *updates, *updates
 
 def remove_custom_lora(selected_indices, current_loras, gallery):
     if current_loras:
@@ -428,30 +372,14 @@ def remove_custom_lora(selected_indices, current_loras, gallery):
     # 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']
+    selected_info, lora_scale, lora_image = get_lora_selected_info(selected_indices, current_loras)
     return (
         current_loras,
         gr.update(value=gallery_items),
-        selected_info_1,
-        selected_info_2,
+        *selected_info,
         selected_indices,
-        lora_scale_1,
-        lora_scale_2,
-        lora_image_1,
-        lora_image_2
+        *lora_scale,
+        *lora_image
     )
 
 @spaces.GPU(duration=70)
@@ -603,7 +531,7 @@ def generate_image_to_image(prompt_mash: str, image_input_path_dict: dict, image
         raise gr.Error(f"I2I Inference Error: {e}") from e
 
 def run_lora(prompt: str, image_input: dict, image_strength: float, task_type: str, turbo_mode: str, blur_mask: bool, blur_factor: float,
-             cfg_scale: float, steps: int, selected_indices, lora_scale_1: float, lora_scale_2: float,
+             cfg_scale: float, steps: int, selected_indices, lora_scale: list[float],
              randomize_seed: bool, seed: int, width: int, height: int, sigmas_factor: float, loras_state,
              lora_json, cn_on: bool, translate_on: bool, progress=gr.Progress(track_tqdm=True)):
     global pipe, pipe_i2i
@@ -665,7 +593,8 @@ def run_lora(prompt: str, image_input: dict, image_strength: float, task_type: s
                 lora_name = f"lora_{idx}"
                 lora_names.append(lora_name)
                 print(f"Lora Name: {lora_name}")
-                lora_weights.append(lora_scale_1 if idx == 0 else lora_scale_2)
+                lora_weights.append(lora_scale[idx])
+                print(f"Lora Weight: {lora_scale[idx]}")
                 lora_path = lora['repo']
                 weight_name = lora.get("weights")
                 print(f"Lora Path: {lora_path}")
@@ -809,7 +738,9 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
             elem_id="title",
         )
         loras_state = gr.State(loras)
+        loras_scale = gr.State([1.15] * MAX_LORA)
         selected_indices = gr.State([])
+        num = [gr.State(value=i) for i in range(MAX_LORA)]
         with gr.Row():
             with gr.Column(scale=3):
                 with gr.Group():
@@ -833,29 +764,24 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
         with gr.Row(elem_id="loaded_loras"):
             with gr.Column(scale=1, min_width=25):
                 randomize_button = gr.Button("🎲", variant="secondary", scale=1, elem_id="random_btn")
-            with gr.Column(scale=8):
-                with gr.Row():
-                    with gr.Column(scale=0, min_width=50):
-                        lora_image_1 = gr.Image(label="LoRA 1 Image", interactive=False, min_width=50, width=50, show_label=False, show_share_button=False, show_download_button=False, show_fullscreen_button=False, height=50)
-                    with gr.Column(scale=3, min_width=100):
-                        selected_info_1 = gr.Markdown("Select a LoRA 1")
-                    with gr.Column(scale=5, min_width=50):
-                        lora_scale_1 = gr.Slider(label="LoRA 1 Scale", minimum=0, maximum=3, step=0.01, value=1.15)
-                with gr.Row():
-                    remove_button_1 = gr.Button("Remove", size="sm")
-            with gr.Column(scale=8):
-                with gr.Row():
-                    with gr.Column(scale=0, min_width=50):
-                        lora_image_2 = gr.Image(label="LoRA 2 Image", interactive=False, min_width=50, width=50, show_label=False, show_share_button=False, show_download_button=False, show_fullscreen_button=False, height=50)
-                    with gr.Column(scale=3, min_width=100):
-                        selected_info_2 = gr.Markdown("Select a LoRA 2")
-                    with gr.Column(scale=5, min_width=50):
-                        lora_scale_2 = gr.Slider(label="LoRA 2 Scale", minimum=0, maximum=3, step=0.01, value=1.15)
-                with gr.Row():
-                    remove_button_2 = gr.Button("Remove", size="sm")
+            lora_image = [None] * MAX_LORA
+            selected_info = [None] * MAX_LORA
+            lora_scale = [None] * MAX_LORA
+            remove_button = [None] * MAX_LORA
+            for i in range(MAX_LORA):
+                with gr.Column(scale=8):
+                    with gr.Row():
+                        with gr.Column(scale=0, min_width=50):
+                            lora_image[i] = gr.Image(label=f"LoRA {int(i+1)} Image", interactive=False, min_width=50, width=50, show_label=False, show_share_button=False, show_download_button=False, show_fullscreen_button=False, height=50)
+                        with gr.Column(scale=3, min_width=100):
+                            selected_info[i] = gr.Markdown(f"Select a LoRA {int(i+1)}")
+                        with gr.Column(scale=5, min_width=50):
+                            lora_scale[i] = gr.Slider(label=f"LoRA {int(i+1)} Scale", minimum=0, maximum=3, step=0.01, value=1.15)
+                    with gr.Row():
+                        remove_button[i] = gr.Button("Remove", size="sm")
         with gr.Row():
             with gr.Column():
-                selected_info = gr.Markdown("")
+                #selected_info = gr.Markdown("")
                 gallery = gr.Gallery([(item["image"], item["title"]) for item in loras], label="LoRA Gallery", allow_preview=False,
                                      columns=4, elem_id="gallery", show_share_button=False, interactive=False)
                 with gr.Group():
@@ -982,31 +908,29 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
     gallery.select(
         update_selection,
         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, 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, loras_state],
-        outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
-    )
+        outputs=[prompt, *selected_info, selected_indices, *lora_scale, width, height, *lora_image])
+    for i in range(MAX_LORA):
+        remove_button[i].click(
+            remove_lora,
+            inputs=[selected_indices, loras_state, num[i]],
+            outputs=[*selected_info, selected_indices, *lora_scale, *lora_image]
+        )
+        lora_scale[i].change(lambda n, l, s: [s if i == n else x for i, x in enumerate(l)], [num[i], loras_scale, lora_scale[i]], [loras_scale])
+
     randomize_button.click(
         randomize_loras,
         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]
+        outputs=[*selected_info, selected_indices, *lora_scale, *lora_image, prompt]
     )
     add_custom_lora_button.click(
         add_custom_lora,
         inputs=[custom_lora, selected_indices, loras_state, gallery],
-        outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
+        outputs=[loras_state, gallery, *selected_info, selected_indices, *lora_scale, *lora_image]
     )
     remove_custom_lora_button.click(
         remove_custom_lora,
         inputs=[selected_indices, loras_state, gallery],
-        outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
+        outputs=[loras_state, gallery, *selected_info, selected_indices, *lora_scale, *lora_image]
     )
     gr.on(
         triggers=[generate_button.click, prompt.submit],
@@ -1018,7 +942,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
         trigger_mode="once",
     ).success(
         fn=run_lora,
-        inputs=[prompt, input_image, image_strength, task_type, turbo_mode, blur_mask, blur_factor, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2,
+        inputs=[prompt, input_image, image_strength, task_type, turbo_mode, blur_mask, blur_factor, cfg_scale, steps, selected_indices, loras_scale,
                 randomize_seed, seed, width, height, sigmas_factor, loras_state, lora_repo_json, cn_on, auto_trans], 
         outputs=[result, seed, progress_bar],
         queue=True,

env.py

@@ -6,7 +6,8 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 HF_READ_TOKEN = os.environ.get('HF_READ_TOKEN') # only use for private repo
 
 
-num_loras = 5
+MAX_LORA = 5 # number of internal LoRA slots
+num_loras = 5 # number of external LoRA slots
 num_cns = 2