v0.018

app.py

@@ -25,7 +25,7 @@ with open('loras.json', 'r') as f:
 
 # Initialize the base model
 dtype = torch.bfloat16
-device = torch.device(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
+device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
 
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
@@ -95,16 +95,49 @@ def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, h
             return gr.update(), "Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6", selected_indices, 1.15, 1.15, 1.15, 1.15, 1.15, 1.15, width, height, None, None, None, None, None, None
 
     # Initialize defaults
-    selected_info = ["Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6"]
-    lora_scale = [1.15] * 6
-    lora_image = [None] * 6
-
-    # Update selected LoRAs
-    for idx, sel_idx in enumerate(selected_indices):
-        if idx < 6:
-            lora = loras_state[sel_idx]
-            selected_info[idx] = f"### LoRA {idx+1} Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
-            lora_image[idx] = lora['image']
+    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"
+    selected_info_5 = "Select a LoRA 5"
+    selected_info_6 = "Select a LoRA 6"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_scale_3 = 1.15
+    lora_scale_4 = 1.15
+    lora_scale_5 = 1.15
+    lora_scale_6 = 1.15
+    lora_image_1 = None
+    lora_image_2 = None
+    lora_image_3 = None
+    lora_image_4 = None
+    lora_image_5 = None
+    lora_image_6 = 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: [{lora3['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 len(selected_indices) >= 5:
+        lora5 = loras_state[selected_indices[4]]
+        selected_info_5 = f"### LoRA 5 Selected: [{lora5['title']}](https://huggingface.co/{lora5['repo']}) ✨"
+        lora_image_5 = lora5['image']
+    if len(selected_indices) >= 6:
+        lora6 = loras_state[selected_indices[5]]
+        selected_info_6 = f"### LoRA 6 Selected: [{lora6['title']}](https://huggingface.co/{lora6['repo']}) ✨"
+        lora_image_6 = lora6['image']
 
     if selected_indices:
         last_selected_lora = loras_state[selected_indices[-1]]
@@ -112,38 +145,196 @@ def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, h
     else:
         new_placeholder = "Type a prompt after selecting a LoRA"
 
-    return gr.update(placeholder=new_placeholder), *selected_info, selected_indices, *lora_scale, width, height, *lora_image
-
-def remove_lora(idx, selected_indices, loras_state):
-    if len(selected_indices) >= idx + 1:
-        selected_indices.pop(idx)
-    selected_info = ["Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6"]
-    lora_scale = [1.15] * 6
-    lora_image = [None] * 6
+    return (
+        gr.update(placeholder=new_placeholder),
+        selected_info_1,
+        selected_info_2,
+        selected_info_3,
+        selected_info_4,
+        selected_info_5,
+        selected_info_6,
+        selected_indices,
+        lora_scale_1,
+        lora_scale_2,
+        lora_scale_3,
+        lora_scale_4,
+        lora_scale_5,
+        lora_scale_6,
+        width,
+        height,
+        lora_image_1,
+        lora_image_2,
+        lora_image_3,
+        lora_image_4,
+        lora_image_5,
+        lora_image_6
+    )
 
-    for idx_sel, sel_idx in enumerate(selected_indices):
-        if idx_sel < 6:
-            lora = loras_state[sel_idx]
-            selected_info[idx_sel] = f"### LoRA {idx_sel+1} Selected: [{lora['title']}]({lora['repo']}) ✨"
-            lora_image[idx_sel] = lora['image']
+def remove_lora_1(selected_indices, loras_state):
+    if len(selected_indices) >= 1:
+        selected_indices.pop(0)
+    return update_selection_after_removal(selected_indices, loras_state)
+
+def remove_lora_2(selected_indices, loras_state):
+    if len(selected_indices) >= 2:
+        selected_indices.pop(1)
+    return update_selection_after_removal(selected_indices, loras_state)
+
+def remove_lora_3(selected_indices, loras_state):
+    if len(selected_indices) >= 3:
+        selected_indices.pop(2)
+    return update_selection_after_removal(selected_indices, loras_state)
+
+def remove_lora_4(selected_indices, loras_state):
+    if len(selected_indices) >= 4:
+        selected_indices.pop(3)
+    return update_selection_after_removal(selected_indices, loras_state)
+
+def remove_lora_5(selected_indices, loras_state):
+    if len(selected_indices) >= 5:
+        selected_indices.pop(4)
+    return update_selection_after_removal(selected_indices, loras_state)
+
+def remove_lora_6(selected_indices, loras_state):
+    if len(selected_indices) >= 6:
+        selected_indices.pop(5)
+    return update_selection_after_removal(selected_indices, loras_state)
+
+def update_selection_after_removal(selected_indices, loras_state):
+    # Reinitialize defaults
+    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"
+    selected_info_5 = "Select a LoRA 5"
+    selected_info_6 = "Select a LoRA 6"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_scale_3 = 1.15
+    lora_scale_4 = 1.15
+    lora_scale_5 = 1.15
+    lora_scale_6 = 1.15
+    lora_image_1 = None
+    lora_image_2 = None
+    lora_image_3 = None
+    lora_image_4 = None
+    lora_image_5 = None
+    lora_image_6 = None
 
-    return *selected_info, selected_indices, *lora_scale, *lora_image
+    # Update selected LoRAs
+    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']
+    if len(selected_indices) >= 3:
+        lora3 = loras_state[selected_indices[2]]
+        selected_info_3 = f"### LoRA 3 Selected: [{lora3['title']}]({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']}]({lora4['repo']}) ✨"
+        lora_image_4 = lora4['image']
+    if len(selected_indices) >= 5:
+        lora5 = loras_state[selected_indices[4]]
+        selected_info_5 = f"### LoRA 5 Selected: [{lora5['title']}]({lora5['repo']}) ✨"
+        lora_image_5 = lora5['image']
+    if len(selected_indices) >= 6:
+        lora6 = loras_state[selected_indices[5]]
+        selected_info_6 = f"### LoRA 6 Selected: [{lora6['title']}]({lora6['repo']}) ✨"
+        lora_image_6 = lora6['image']
+
+    return (
+        selected_info_1,
+        selected_info_2,
+        selected_info_3,
+        selected_info_4,
+        selected_info_5,
+        selected_info_6,
+        selected_indices,
+        lora_scale_1,
+        lora_scale_2,
+        lora_scale_3,
+        lora_scale_4,
+        lora_scale_5,
+        lora_scale_6,
+        lora_image_1,
+        lora_image_2,
+        lora_image_3,
+        lora_image_4,
+        lora_image_5,
+        lora_image_6
+    )
 
 def randomize_loras(selected_indices, loras_state):
     num_loras = min(6, len(loras_state))
     selected_indices = random.sample(range(len(loras_state)), num_loras)
-    selected_info = ["Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6"]
-    lora_scale = [1.15] * 6
-    lora_image = [None] * 6
+    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"
+    selected_info_5 = "Select a LoRA 5"
+    selected_info_6 = "Select a LoRA 6"
+    lora_scale_1 = 1.15
+    lora_scale_2 = 1.15
+    lora_scale_3 = 1.15
+    lora_scale_4 = 1.15
+    lora_scale_5 = 1.15
+    lora_scale_6 = 1.15
+    lora_image_1 = None
+    lora_image_2 = None
+    lora_image_3 = None
+    lora_image_4 = None
+    lora_image_5 = None
+    lora_image_6 = None
 
     for idx, sel_idx in enumerate(selected_indices):
-        if idx < 6:
-            lora = loras_state[sel_idx]
-            selected_info[idx] = f"### LoRA {idx+1} Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
-            lora_image[idx] = lora['image']
+        lora = loras_state[sel_idx]
+        if idx == 0:
+            selected_info_1 = f"### LoRA 1 Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
+            lora_image_1 = lora['image']
+        elif idx == 1:
+            selected_info_2 = f"### LoRA 2 Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
+            lora_image_2 = lora['image']
+        elif idx == 2:
+            selected_info_3 = f"### LoRA 3 Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
+            lora_image_3 = lora['image']
+        elif idx == 3:
+            selected_info_4 = f"### LoRA 4 Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
+            lora_image_4 = lora['image']
+        elif idx == 4:
+            selected_info_5 = f"### LoRA 5 Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
+            lora_image_5 = lora['image']
+        elif idx == 5:
+            selected_info_6 = f"### LoRA 6 Selected: [{lora['title']}](https://huggingface.co/{lora['repo']}) ✨"
+            lora_image_6 = lora['image']
 
     random_prompt = random.choice(prompt_values)
-    return *selected_info, selected_indices, *lora_scale, *lora_image, random_prompt
+    return (
+        selected_info_1,
+        selected_info_2,
+        selected_info_3,
+        selected_info_4,
+        selected_info_5,
+        selected_info_6,
+        selected_indices,
+        lora_scale_1,
+        lora_scale_2,
+        lora_scale_3,
+        lora_scale_4,
+        lora_scale_5,
+        lora_scale_6,
+        lora_image_1,
+        lora_image_2,
+        lora_image_3,
+        lora_image_4,
+        lora_image_5,
+        lora_image_6,
+        random_prompt
+    )
 
 def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
     if custom_lora:
@@ -174,36 +365,19 @@ def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
                 gr.Warning("You can select up to 6 LoRAs, remove one to select a new one.")
 
             # Update selected_info and images
-            selected_info = ["Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6"]
-            lora_scale = [1.15] * 6
-            lora_image = [None] * 6
-
-            for idx, sel_idx in enumerate(selected_indices):
-                if idx < 6:
-                    lora = current_loras[sel_idx]
-                    selected_info[idx] = f"### LoRA {idx+1} Selected: {lora['title']} ✨"
-                    lora_image[idx] = lora['image'] if lora['image'] else None
-
-            print("Finished adding custom LoRA")
-            return (
-                current_loras,
-                gr.update(value=gallery_items),
-                *selected_info,
-                selected_indices,
-                *lora_scale,
-                *lora_image
-            )
+            return update_selection_after_removal(selected_indices, current_loras)[:20] + (current_loras, gr.update(value=gallery_items))
+
         except Exception as e:
             print(e)
             gr.Warning(str(e))
-            return current_loras, gr.update(), *["Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6"], selected_indices, *[1.15]*6, *[None]*6
+            return current_loras, gr.update(), "Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6", selected_indices, 1.15, 1.15, 1.15, 1.15, 1.15, 1.15, None, None, None, None, None, None
     else:
-        return current_loras, gr.update(), *["Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6"], selected_indices, *[1.15]*6, *[None]*6
+        return current_loras, gr.update(), "Select a LoRA 1", "Select a LoRA 2", "Select a LoRA 3", "Select a LoRA 4", "Select a LoRA 5", "Select a LoRA 6", selected_indices, 1.15, 1.15, 1.15, 1.15, 1.15, 1.15, None, None, None, None, None, None
 
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
     print("Generating image...")
-    pipe.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
-    generator = torch.Generator(device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')).manual_seed(seed)
+    pipe.to("cuda")
+    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(
@@ -220,8 +394,8 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             yield img
 
 def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, seed):
-    pipe_i2i.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
-    generator = torch.Generator(device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')).manual_seed(seed)
+    pipe_i2i.to("cuda")
+    generator = torch.Generator(device="cuda").manual_seed(seed)
     image_input = load_image(image_input_path)
     final_image = pipe_i2i(
         prompt=prompt_mash,
@@ -238,7 +412,26 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     return final_image
 
 @spaces.GPU(duration=75)
-def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, lora_scale_5, lora_scale_6, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
+def run_lora(
+    prompt,
+    image_input,
+    image_strength,
+    cfg_scale,
+    steps,
+    selected_indices,
+    lora_scale_1,
+    lora_scale_2,
+    lora_scale_3,
+    lora_scale_4,
+    lora_scale_5,
+    lora_scale_6,
+    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.")
 
@@ -265,7 +458,14 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     # Load LoRA weights with respective scales
     lora_names = []
     lora_weights = []
-    lora_scales = [lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, lora_scale_5, lora_scale_6]
+    lora_scales = [
+        lora_scale_1,
+        lora_scale_2,
+        lora_scale_3,
+        lora_scale_4,
+        lora_scale_5,
+        lora_scale_6
+    ]
     with calculateDuration("Loading LoRA weights"):
         for idx, lora in enumerate(selected_loras):
             if idx >= 6:
@@ -313,7 +513,7 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
         yield final_image, seed, gr.update(value=progress_bar, visible=False)
 
 run_lora.zerogpu = True
-
+@spaces.GPU
 def get_huggingface_safetensors(link):
     split_link = link.split("/")
     if len(split_link) == 2: