Update app.py

app.py

@@ -89,144 +89,70 @@ def calculate_optimal_dimensions(image):
     return width, height
 
 
-@spaces.GPU(duration=30)
-def inpaint(
-        image,
-        mask,
-        prompt="",
-        seed=0,
-        num_inference_steps=28,
-        guidance_scale=50,
-        flux_keywords: List[str] = None,
-        loras: List[LoRA] = None
-):
-    """
-    Runs inpainting with selected LoRAs and their keywords.
-    """
-
-    # Step 1: Reset LoRAs
-    deactivate_loras(pipe)
-
-    # Step 2: Prepare selected LoRAs and load them
-    selected_loras = {}
-    for lora in loras:
-        if lora.url:
-            selected_loras[lora.url] = 1.0  # Default weight, could be made configurable
-
-    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 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,
-        width=width,
-        height=height,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        generator=torch.Generator().manual_seed(seed)
-    ).images[0]
-
-    return result.convert("RGBA"), final_prompt, seed
-
-
-def process_ui_inputs(*ui_args):
+def parse_ui_inputs(*ui_args):
     """
-    Processes UI inputs and converts them to the format expected by inpaint function.
+    Parses UI inputs and extracts flux keywords and LoRAs with weights.
+    Returns: (main_inputs, selected_flux_keywords, selected_loras_with_weights)
     """
-    # Extract main inputs
-    image = ui_args[0]
-    mask = ui_args[1]
-    prompt = ui_args[2]
-    seed = ui_args[3]
-    num_inference_steps = ui_args[4]
-    guidance_scale = ui_args[5]
-
+    # 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 = []
+    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:
-                # Create a copy of the LoRA with updated weight (stored in a custom way)
-                lora_copy = LoRA(
-                    id=lora_config.id,
-                    nombre=lora_config.nombre,
-                    title=lora_config.title,
-                    url=lora_config.url,
-                    Keywords=lora_config.Keywords,
-                    note=lora_config.note
-                )
-                # Store weight in a way that can be used by activate_loras
-                selected_loras.append((lora_copy, weight))
-
-    return inpaint_with_weights(
-        image, mask, prompt, seed, num_inference_steps, guidance_scale,
-        selected_flux_keywords, selected_loras
-    )
+                selected_loras_with_weights.append((lora_config, weight))
+    
+    return main_inputs, selected_flux_keywords, selected_loras_with_weights
 
 
-def inpaint_with_weights(
-        image, mask, prompt, seed, num_inference_steps, guidance_scale,
-        flux_keywords, loras_with_weights
+@spaces.GPU(duration=30)
+def inpaint(
+        image,
+        mask,
+        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
 ):
     """
-    Modified inpaint function that handles LoRAs with weights.
+    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)
@@ -277,6 +203,26 @@ def inpaint_with_weights(
     return result.convert("RGBA"), final_prompt, seed
 
 
+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,
+        prompt=prompt,
+        seed=seed,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        flux_keywords=selected_flux_keywords,
+        loras_with_weights=selected_loras_with_weights
+    )
+
+
 def toggle_input(checked, current_value):
     """
     Enables or disables the Number input based on checkbox status.
@@ -347,7 +293,7 @@ def create_lora_components():
 
 
 # Create main interface
-with gr.Blocks(title="Flux.1 Fill dev Inpainting with LoRAs from JSON", theme=gr.themes.Soft()) as demo:
+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):
             prompt_input = gr.Text(
@@ -402,7 +348,7 @@ with gr.Blocks(title="Flux.1 Fill dev Inpainting with LoRAs from JSON", theme=gr
 
     if all_components:
         run_btn.click(
-            process_ui_inputs,
+            inpaint_ui_wrapper,
             inputs=[
                 image_input,
                 mask_input,