Update app.py

app.py

@@ -1,7 +1,4 @@
-import json
-import os
 import random
-from dataclasses import field
 from typing import List
 
 import gradio as gr
@@ -9,7 +6,6 @@ import numpy as np
 import spaces
 import torch
 from diffusers import FluxFillPipeline
-
 from loras import LoRA, loras
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -18,7 +14,6 @@ pipe = FluxFillPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
 ).to("cuda")
 
-
 # Flux system keywords list
 flux_keywords_available = [
     "IMG_1025.HEIC",
@@ -32,13 +27,14 @@ flux_keywords_available = [
 ]
 
 
-def activate_loras(pipe, loras_with_weights):
+def activate_loras(pipe: FluxFillPipeline, loras_with_weights: list[tuple[LoRA, float]]):
     """
     Loads the selected LoRAs into the pipeline with the specified weights.
     """
-    for repo, weight in loras_with_weights.items():
-        print(f"Loading LoRA: {repo} with weight {weight}")
-        pipe.load_lora_weights(repo, weight=weight)
+    for lora, weight in loras_with_weights:
+        print(f"Loading LoRA: {lora.display_name} with weight {weight}")
+        pipe.load_lora_weights(lora.id, weight=weight)
+
     return pipe
 
 
@@ -51,6 +47,53 @@ def deactivate_loras(pipe):
     return pipe
 
 
+def inpaint_wrapper(
+        image,
+        mask,
+        prompt="",
+        seed=0,
+        num_inference_steps=28,
+        guidance_scale=50,
+        flux_keywords: List[str] = None,
+        loras_with_weights: List[tuple[LoRA, float]] = None  # List of (LoRA, weight) tuples
+):
+
+    # Step 1: Reset LoRAs
+    deactivate_loras(pipe)
+
+    # Step 2: Prepare selected LoRAs and load them
+    if loras_with_weights:
+        activate_loras(pipe, loras_with_weights)
+
+    # Step 3: Prepare prompt
+    final_prompt = ""
+
+    for keyword in flux_keywords:
+        final_prompt += f"{keyword}, "
+
+    for lora, _ in loras_with_weights:
+        if lora.Keyword:
+            keywords_str = ", ".join(lora.Keywords)
+            final_prompt += f"{keywords_str}, "
+
+    if prompt:
+        final_prompt += "\n\n"
+
+    final_prompt += prompt
+
+    # Step 4: Seed handling
+    if not isinstance(seed, int) or seed < 0:
+        seed = random.randint(0, MAX_SEED)
+
+    return inpaint(
+        image,
+        mask,
+        prompt=final_prompt,
+        seed=seed,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+    )
+
 def calculate_optimal_dimensions(image):
     # Extract original dimensions
     original_width, original_height = image.size
@@ -89,44 +132,6 @@ def calculate_optimal_dimensions(image):
     return width, height
 
 
-def parse_ui_inputs(*ui_args):
-    """
-    Parses UI inputs and extracts flux keywords and LoRAs with weights.
-    Returns: (main_inputs, selected_flux_keywords, selected_loras_with_weights)
-    """
-    # Extract main inputs (first 6 arguments)
-    main_inputs = ui_args[:6]  # image, mask, prompt, seed, steps, guidance
-    
-    # Extract flux keyword selections
-    num_flux_keywords = len(flux_keywords_available)
-    flux_keyword_selections = ui_args[6:6 + num_flux_keywords]
-    
-    # Extract LoRA selections and weights
-    lora_args_start = 6 + num_flux_keywords
-    lora_args = ui_args[lora_args_start:]
-    
-    # Process selected flux keywords
-    selected_flux_keywords = []
-    for i, selected in enumerate(flux_keyword_selections):
-        if selected:
-            selected_flux_keywords.append(flux_keywords_available[i])
-    
-    # Process selected LoRAs with weights
-    selected_loras_with_weights = []
-    for i, lora_config in enumerate(loras):
-        checkbox_idx = i * 2
-        weight_idx = i * 2 + 1
-        
-        if checkbox_idx < len(lora_args):
-            checked = lora_args[checkbox_idx]
-            weight = lora_args[weight_idx] if weight_idx < len(lora_args) else 0.5
-            
-            if checked:
-                selected_loras_with_weights.append((lora_config, weight))
-    
-    return main_inputs, selected_flux_keywords, selected_loras_with_weights
-
-
 @spaces.GPU(duration=30)
 def inpaint(
         image,
@@ -134,65 +139,22 @@ def inpaint(
         prompt="",
         seed=0,
         num_inference_steps=28,
-        guidance_scale=50,
-        flux_keywords: List[str] = None,
-        loras_with_weights: List[tuple] = None  # List of (LoRA, weight) tuples
+        guidance_scale=50
 ):
     """
     Main inpainting function with selected LoRAs and their keywords.
     """
-    if flux_keywords is None:
-        flux_keywords = []
-    if loras_with_weights is None:
-        loras_with_weights = []
-
-    # Step 1: Reset LoRAs
-    deactivate_loras(pipe)
-
-    # Step 2: Prepare selected LoRAs and load them
-    selected_loras = {}
-    active_loras = []
-    
-    for lora, weight in loras_with_weights:
-        if lora.url:
-            selected_loras[lora.url] = round(weight, 1)
-            active_loras.append(lora)
-
-    if selected_loras:
-        activate_loras(pipe, selected_loras)
 
     print("ACTIVE ADAPTERS:", pipe.get_active_adapters())
 
-    # Step 3: Prepare prompt
     image = image.convert("RGB")
     mask = mask.convert("L")
     width, height = calculate_optimal_dimensions(image)
 
-    final_prompt = ""
-
-    # Add selected flux keywords
-    for keyword in flux_keywords:
-        final_prompt += f"{keyword}, "
-
-    # Add keywords from active LoRAs
-    for lora in active_loras:
-        if lora.Keywords:
-            keywords_str = ", ".join(lora.Keywords)
-            final_prompt += f"{keywords_str}, "
-
-    if prompt:
-        final_prompt += "\n\n"
-    final_prompt += prompt
-
-    # Step 4: Seed handling
-    if not isinstance(seed, int) or seed <= 0:
-        seed = random.randint(0, MAX_SEED)
-
-    # Step 5: Run pipeline
     result = pipe(
         image=image,
         mask_image=mask,
-        prompt=final_prompt,
+        prompt=prompt,
         width=width,
         height=height,
         num_inference_steps=num_inference_steps,
@@ -200,17 +162,54 @@ def inpaint(
         generator=torch.Generator().manual_seed(seed)
     ).images[0]
 
-    return result.convert("RGBA"), final_prompt, seed
+    return result.convert("RGBA"), prompt, seed
+
+'''
+def parse_ui_inputs(*ui_args):
+    """
+    Parses UI inputs and extracts flux keywords and LoRAs with weights.
+    Returns: (main_inputs, selected_flux_keywords, selected_loras_with_weights)
+    """
+    # Extract main inputs (first 6 arguments)
+    main_inputs = ui_args[:6]  # image, mask, prompt, seed, steps, guidance
+
+    # Extract flux keyword selections
+    num_flux_keywords = len(flux_keywords_available)
+    flux_keyword_selections = ui_args[6:6 + num_flux_keywords]
+
+    # Extract LoRA selections and weights
+    lora_args_start = 6 + num_flux_keywords
+    lora_args = ui_args[lora_args_start:]
+
+    # Process selected flux keywords
+    selected_flux_keywords = []
+    for i, selected in enumerate(flux_keyword_selections):
+        if selected:
+            selected_flux_keywords.append(flux_keywords_available[i])
+
+    # Process selected LoRAs with weights
+    selected_loras_with_weights = []
+    for i, lora_config in enumerate(loras):
+        checkbox_idx = i * 2
+        weight_idx = i * 2 + 1
+
+        if checkbox_idx < len(lora_args):
+            checked = lora_args[checkbox_idx]
+            weight = lora_args[weight_idx] if weight_idx < len(lora_args) else 0.5
 
+            if checked:
+                selected_loras_with_weights.append((lora_config, weight))
+
+    return main_inputs, selected_flux_keywords, selected_loras_with_weights
 
 def inpaint_ui_wrapper(*ui_args):
     """
     UI wrapper that processes Gradio inputs and calls the main inpaint function.
     """
     main_inputs, selected_flux_keywords, selected_loras_with_weights = parse_ui_inputs(*ui_args)
-    
+
     image, mask, prompt, seed, num_inference_steps, guidance_scale = main_inputs
-    
+
     return inpaint(
         image=image,
         mask=mask,
@@ -221,8 +220,27 @@ def inpaint_ui_wrapper(*ui_args):
         flux_keywords=selected_flux_keywords,
         loras_with_weights=selected_loras_with_weights
     )
+'''
+def selected_keywords(keyword_components):
+
+    keywords = []
+    for checkbox in keyword_components:
+        if checkbox.value:
+            keywords.append(checkbox.label)
+
+    return keywords
 
+def selected_loras(components: list[tuple[gr.Checkbox, gr.Number]]) -> list[tuple[LoRA, float]]:
+
+    selected_loras_with_weights = []
+    for checkbox, number in components:
+        for lora in loras:
+            if checkbox.value and lora.display_name == checkbox.label:
+                selected_loras_with_weights.append((lora, number.value))
 
+    return selected_loras_with_weights
+
+'''
 def toggle_input(checked, current_value):
     """
     Enables or disables the Number input based on checkbox status.
@@ -235,7 +253,7 @@ def toggle_input(checked, current_value):
         interactive=checked,
         label="Adapter weight"
     )
-
+'''
 
 def create_flux_keyword_components():
     """
@@ -257,14 +275,14 @@ def create_lora_components():
     """
     Creates interface components for each LoRA from the imported loras list.
     """
+
     components = []
 
     for lora in loras:
         with gr.Row():
             # Create checkbox with LoRA title
-            display_name = lora.title or lora.nombre or "Unnamed LoRA"
             checkbox = gr.Checkbox(
-                label=display_name,
+                label=lora.display_name,
                 value=False,
                 info=lora.note or ""  # Show note as additional info
             )
@@ -278,7 +296,7 @@ def create_lora_components():
                 interactive=False,
                 label="Adapter weight"
             )
-
+            """
             # Connect checkbox with number input
             checkbox.change(
                 toggle_input,
@@ -286,13 +304,12 @@ def create_lora_components():
                 outputs=number_input,
                 api_name=False
             )
-
-            components.extend([checkbox, number_input])
+            """
+            components.append((checkbox, number_input))
 
     return components
 
 
-# Create main interface
 with gr.Blocks(title="Flux.1 Fill dev Inpainting with LoRAs", theme=gr.themes.Soft()) as demo:
     with gr.Row():
         with gr.Column(scale=2):
@@ -326,12 +343,10 @@ with gr.Blocks(title="Flux.1 Fill dev Inpainting with LoRAs", theme=gr.themes.So
             flux_keyword_components = create_flux_keyword_components()
 
             # LoRAs section
-            gr.Markdown("### Available LoRAs")
+            lora_components = []
             if loras:
+                gr.Markdown("### Available LoRAs")
                 lora_components = create_lora_components()
-            else:
-                gr.Markdown("*No LoRAs found*")
-                lora_components = []
 
         with gr.Column(scale=3):
             run_btn = gr.Button("Run", variant="primary")
@@ -343,37 +358,22 @@ with gr.Blocks(title="Flux.1 Fill dev Inpainting with LoRAs", theme=gr.themes.So
             used_prompt_box = gr.Text(label="Used prompt", lines=4)
             used_seed_box = gr.Number(label="Used seed")
 
-    # Connect button with function
-    all_components = flux_keyword_components + lora_components
-
-    if all_components:
-        run_btn.click(
-            inpaint_ui_wrapper,
-            inputs=[
-                image_input,
-                mask_input,
-                prompt_input,
-                seed_slider,
-                num_inference_steps_input,
-                guidance_scale_input,
-                *all_components
-            ],
-            outputs=[result_image, used_prompt_box, used_seed_box]
-        )
-    else:
-        # If no components, create simplified version
-        run_btn.click(
-            lambda *args: (None, "No LoRAs configured", 0),
-            inputs=[
-                image_input,
-                mask_input,
-                prompt_input,
-                seed_slider,
-                num_inference_steps_input,
-                guidance_scale_input
-            ],
-            outputs=[result_image, used_prompt_box, used_seed_box]
-        )
+    run_btn.click(
+        inpaint_wrapper,
+        inputs=[
+            image_input,
+            mask_input,
+            prompt_input,
+            seed_slider,
+            num_inference_steps_input,
+            guidance_scale_input,
+            selected_keywords(flux_keyword_components),
+            selected_loras(lora_components)
+        ],
+        outputs=[result_image, used_prompt_box, used_seed_box],
+        api_name="/inpaint"
+    )
+
 
 if __name__ == "__main__":
-    demo.launch(share=False, show_error=True)
+    demo.launch(share=False, show_error=True)