Update app.py

app.py

@@ -92,17 +92,56 @@ def download_file(url, directory=None):
         file.write(response.content)
     
     return filepath
-
-def get_trigger_word(base_model, lora_models):
-    trigger_words = []  # Initialize an empty list to hold trigger words
-    for lora in lora_models:
-        trigger_word = lora.get('trigger_word', '')
-        if trigger_word:
-            trigger_words.append(trigger_word)
     
-    # Join all trigger words with a space and wrap in a span
-    trigger_word = "<span style='color:green;'> " + " ".join(trigger_words) + " </span>"
-    return trigger_word
+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:
+        selected_indices.remove(selected_index)
+    else:
+        if len(selected_indices) < 4:
+            selected_indices.append(selected_index)
+        else:
+            gr.Warning("You can select up to 4 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(), gr.update(), gr.update()
+
+    selected_info_1 = "Select a LoRA 1"
+    selected_info_2 = "Select a LoRA 2"
+    selected_info_3 = "Select a LoRA 3"
+    selected_info_4 = "Select a LoRA 4"
+    lora_scale_1 = 0.5
+    lora_scale_2 = 0.5
+    lora_scale_3 = 0.5
+    lora_scale_4 = 0.5
+    lora_image_1 = None
+    lora_image_2 = None
+    lora_image_3 = None
+    lora_image_4 = 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']
+    if len(selected_indices) >= 3:
+        lora3 = loras_state[selected_indices[2]]
+        selected_info_3 = f"### LoRA 3 Selected: [{lora2['title']}](https://huggingface.co/{lora3['repo']}) ✨"
+        lora_image_3 = lora3['image']
+    if len(selected_indices) >= 4:
+        lora4 = loras_state[selected_indices[3]]
+        selected_info_4 = f"### LoRA 4 Selected: [{lora4['title']}](https://huggingface.co/{lora4['repo']}) ✨"
+        lora_image_4 = lora4['image']
+
+    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 after selecting a LoRA"
+
+    return gr.update(placeholder=new_placeholder), selected_info_1, selected_info_2, selected_info_3, selected_info_4, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, width, height, lora_image_1, lora_image_2, lora_image_3, lora_image_4
+
     
 def randomize_loras(selected_indices, loras_state):
     if len(loras_state) < 2:
@@ -127,42 +166,6 @@ def randomize_loras(selected_indices, loras_state):
     
     return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, random_prompt
 
-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:
-        selected_indices.remove(selected_index)
-    else:
-        if len(selected_indices) < 4:
-            selected_indices.append(selected_index)
-        else:
-            gr.Warning("You can select up to 4 LoRAs, remove one to select a new one.")
-            return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update(), width, height, gr.update(), gr.update(), gr.update()
-
-    selected_info = ["Select a LoRA {}".format(i) for i in range(1, 5)]
-    lora_images = [None] * 4
-    lora_scales = [0.5] * 4
-
-    for i in range(len(selected_indices)):
-        lora = loras_state[selected_indices[i]]
-        trigger_word = lora.get('trigger_word', '')
-        selected_info[i] = f"### LoRA {i + 1} Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨ {trigger_word}"
-        lora_images[i] = lora['image']
-
-    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 after selecting a LoRA"
-
-    return (gr.update(placeholder=new_placeholder), 
-            *selected_info, selected_indices,
-            *lora_scales,
-            width, height, 
-            *lora_images,
-            gr.update()
-    )
 def remove_lora_1(selected_indices, loras_state):
     if len(selected_indices) >= 1:
         selected_indices.pop(0)
@@ -310,24 +313,7 @@ def remove_lora_4(selected_indices, loras_state):
         selected_info_4 = f"### LoRA 4 Selected: [{lora4['title']}]({lora4['repo']}) ✨ {trigger_word}"
         lora_image_4 = lora4['image']
     return selected_info_1, selected_info_2, selected_info_3, selected_info_4, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, lora_image_1, lora_image_2, lora_image_3, lora_image_4, gr.update()
-    
-def remove_lora(selected_indices, loras_state):
-    # Remove the LoRA based on the index
-    if selected_indices:
-        selected_indices.pop()  # Remove the last selected LoRA
-
-    selected_info = ["Select a LoRA"] * 4
-    lora_images = [None] * 4
-    lora_scales = [0.5] * 4
-
-    for i in range(min(len(selected_indices), 4)):
-        lora = loras_state[selected_indices[i]]
-        trigger_word = lora.get('trigger_word', '')
-        selected_info[i] = f"### LoRA {i + 1} Selected: [{lora['title']}]({lora['repo']}) ✨ {trigger_word}"
-        lora_images[i] = lora['image']
-
-    return selected_info, selected_indices, lora_scales, lora_images, gr.update()
-    
+       
 def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
     if custom_lora:
         try:
@@ -395,9 +381,9 @@ def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
         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(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), 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()
+        return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
 
 def remove_custom_lora(selected_indices, current_loras, gallery):
     if current_loras:
@@ -465,13 +451,13 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             output_type="pil",
             good_vae=good_vae,
         ):
-            print("Image generated successfully.")  # Debugging statement
+            #print("Image generated successfully.")  # Debugging statement
             yield img
             final_image = img  # Update final_image with the current image
     return final_image
 
 @spaces.GPU(duration=75)
-def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, randomize_seed, seed, width, height, loras_state, image_input=None, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, cfg_scale, steps, selected_info_1, selected_info_2, selected_info_3, selected_info_4, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, randomize_seed, seed, width, height, loras_state, image_input=None, progress=gr.Progress(track_tqdm=True)):
     print("run_lora function called.")  # Debugging statement
     print(f"Inputs received - Prompt: {prompt}, CFG Scale: {cfg_scale}, Steps: {steps}, Seed: {seed}, Width: {width}, Height: {height}")  # Debugging statement
 
@@ -487,7 +473,7 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
     for lora in selected_loras:
         trigger_word = lora.get('trigger_word', '')
         if trigger_word:
-            if lora.get("trigger_position") == "prepend":
+            if lora.get("trigger_position") == "append":
                 prepends.append(trigger_word)
             else:
                 appends.append(trigger_word)
@@ -540,7 +526,7 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
     try:
         if image_input is not None:
             final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
-            yield final_image, seed, gr.update(visible=True)
+            yield final_image, seed, gr.update(visible=False)
         else:
             image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             # Consume the generator to get the final image
@@ -563,7 +549,7 @@ run_lora.zerogpu = True
 
 def get_huggingface_safetensors(link):
     split_link = link.split("/")
-    if len(split_link) == 2:
+    if len(split_link) == 4:
         model_card = ModelCard.load(link)
         base_model = model_card.data.get("base_model")
         print(f"Base model: {base_model}")