Update app.py

app.py

@@ -160,21 +160,21 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     return final_image 
 
 @spaces.GPU(duration=70)
-def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
-    print("Selected indices (raw):", selected_indices)
-    # Then parse them if needed
-    parsed_indices = parse_selected_indices(selected_indices)
-    print("Parsed indices:", parsed_indices)
+def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices_json, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+    import json
+    # Parse the JSON string into a list of indices
+    selected_indices = json.loads(selected_indices_json)
+    
     # Ensure at least one LoRA is selected
     if not selected_indices or len(selected_indices) == 0:
         raise gr.Error("You must select at least one LoRA before proceeding.")
     
-    # Combine trigger words from all selected LoRAs
+    # Combine trigger words from all selected LoRAs into the prompt
     prompt_mash = prompt
     for idx in selected_indices:
         selected_lora = loras[idx]
         if "trigger_word" in selected_lora and selected_lora["trigger_word"]:
-            # Prepend each trigger word to the prompt
+            # Prepend each trigger word to the prompt; you can adjust the order or separator as needed
             prompt_mash = f"{selected_lora['trigger_word']} {prompt_mash}"
     
     # Unload any previously loaded LoRA weights
@@ -184,14 +184,14 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     
     # Load each selected LoRA weight sequentially
     with calculateDuration("Loading LoRA weights"):
-        # Choose which pipeline to use based on whether an input image is provided
+        # Use the image-to-image pipeline if an input image is provided, else the text-to-image pipeline
         pipe_to_use = pipe_i2i if image_input is not None else pipe
         for idx in selected_indices:
             selected_lora = loras[idx]
             weight_name = selected_lora.get("weights", None)
             pipe_to_use.load_lora_weights(
-                selected_lora["repo"],
-                weight_name=weight_name,
+                selected_lora["repo"], 
+                weight_name=weight_name, 
                 low_cpu_mem_usage=True
             )
     
@@ -202,10 +202,16 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     
     # Generate image(s)
     if image_input is not None:
-        final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
+        # Image-to-image generation
+        final_image = generate_image_to_image(
+            prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed
+        )
         yield final_image, seed, gr.update(visible=False)
     else:
-        image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
+        # Text-to-image generation
+        image_generator = generate_image(
+            prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress
+        )
         final_image = None
         step_counter = 0
         for image in image_generator:
@@ -215,6 +221,7 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             yield image, seed, gr.update(value=progress_bar, visible=True)
         yield final_image, seed, gr.update(value=progress_bar, visible=False)
 
+
         
 def get_huggingface_safetensors(link):
   split_link = link.split("/")
@@ -310,38 +317,40 @@ css = '''
 .progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
 '''
 font=[gr.themes.GoogleFont("Source Sans Pro"), "Arial", "sans-serif"]
+# Begin Gradio Blocks
 with gr.Blocks(theme=gr.themes.Soft(font=font), css=css, delete_cache=(60, 60)) as app:
     title = gr.HTML(
         """<h1><img src="https://huggingface.co/spaces/kayte0342/test/resolve/main/DA4BE61E-A0BD-4254-A1B6-AD3C05D18A9C%20(1).png?download=true" alt="LoRA"> FLUX LoRA Kayte's Space</h1>""",
         elem_id="title",
     )
-    selected_index = gr.State(None)
+    
+    # Hidden textbox to store the JSON string of selected indices
+    selected_indices_hidden = gr.Textbox(value="[]", visible=False)
+    
     with gr.Row():
         with gr.Column(scale=3):
             prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Type a prompt after selecting a LoRA")
         with gr.Column(scale=1, elem_id="gen_column"):
             generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
+    
     with gr.Row():
         with gr.Column():
             selected_info = gr.Markdown("")
-            gallery = gr.Gallery(
-                [(item["image"], item["title"]) for item in loras],
-                label="LoRA Gallery",
-                #select_mode="multiple",  # Ensure multi-selection is enabled
-                allow_preview=False,
-                columns=3,
-                elem_id="gallery",
-                show_share_button=False
-            )
-            with gr.Group():
-                custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path", placeholder="multimodalart/vintage-ads-flux")
-                gr.Markdown("[Check the list of FLUX LoRas](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
-            custom_lora_info = gr.HTML(visible=False)
-            custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
+            # Create a custom container for LoRA selection with checkboxes
+            lora_selection_container = gr.Column()
+            # We'll collect individual checkbox components in a list for later use
+            lora_checkbox_list = []
+            for idx, lora in enumerate(loras):
+                with gr.Row():
+                    gr.Image(value=lora["image"], label=lora["title"], height=100)
+                    checkbox = gr.Checkbox(label="Select", value=False, elem_id=f"lora_checkbox_{idx}")
+                    lora_checkbox_list.append(checkbox)
+            # Add a hidden update button (invisible) to update the hidden state; it can be triggered programmatically.
+            update_selection_btn = gr.Button("Update LoRA Selection", visible=False)
         with gr.Column():
-            progress_bar = gr.Markdown(elem_id="progress",visible=False)
+            progress_bar = gr.Markdown(elem_id="progress", visible=False)
             result = gr.Image(label="Generated Image")
-
+    
     with gr.Row():
         with gr.Accordion("Advanced Settings", open=False):
             with gr.Row():
@@ -351,34 +360,52 @@ with gr.Blocks(theme=gr.themes.Soft(font=font), css=css, delete_cache=(60, 60))
                 with gr.Row():
                     cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
                     steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
-                
                 with gr.Row():
                     width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
                     height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
-                
                 with gr.Row():
                     randomize_seed = gr.Checkbox(True, label="Randomize seed")
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=2**32-1, step=1, value=0, randomize=True)
                     lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
-
-    gallery.select(
-        update_selection,
-        inputs=[width, height],
-        outputs=[prompt, selected_info, selected_index, width, height]
+    
+    # Function to combine checkbox states into a JSON list of selected indices.
+    def combine_selections(*checkbox_values):
+        selected_indices = [i for i, v in enumerate(checkbox_values) if v]
+        return json.dumps(selected_indices)
+    
+    # When the (invisible) update button is clicked, update the hidden state.
+    update_selection_btn.click(
+        combine_selections,
+        inputs=lora_checkbox_list,
+        outputs=selected_indices_hidden
     )
-    custom_lora.input(
-        add_custom_lora,
-        inputs=[custom_lora],
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
+    
+    # Also update the selected_info Markdown when the hidden state changes.
+    def update_info(selected_json):
+        selected_indices = json.loads(selected_json)
+        if selected_indices:
+            info = "Selected LoRAs: " + ", ".join([loras[i]["title"] for i in selected_indices])
+        else:
+            info = "No LoRAs selected."
+        return info
+    selected_info.change(
+        update_info,
+        inputs=selected_indices_hidden,
+        outputs=selected_info
     )
-    custom_lora_button.click(
-        remove_custom_lora,
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
+    
+    # Also, when the Generate button is clicked, update the hidden state from the checkboxes.
+    generate_button.click(
+        combine_selections,
+        inputs=lora_checkbox_list,
+        outputs=selected_indices_hidden
     )
+    
+    # Finally, trigger the generation function (run_lora). Note that run_lora should be modified to parse the JSON string.
     gr.on(
         triggers=[generate_button.click, prompt.submit],
-        fn=run_lora,
-        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
+        fn=run_lora,  # Make sure run_lora begins by parsing the JSON: selected_indices = json.loads(selected_indices_hidden)
+        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices_hidden, randomize_seed, seed, width, height, lora_scale],
         outputs=[result, seed, progress_bar]
     )