app.py

import random
from typing import List, Tuple

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import FluxFillPipeline
from loras import LoRA, loras

MAX_SEED = np.iinfo(np.int32).max

pipe = FluxFillPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Fill-dev",
    torch_dtype=torch.bfloat16
)

flux_keywords_available = [
    "IMG_1025.HEIC",
    "Selfie",
]


def activate_loras(pipe: FluxFillPipeline, loras_with_weights: list[tuple[LoRA, float]]):
    adapter_names = []
    adapter_weights = []

    for lora, weight in loras_with_weights:
        print(f"Loading LoRA: {lora.name} with weight {weight}")
        pipe.load_lora_weights(lora.id, weight=weight, adapter_name=lora.name)
        adapter_names.append(lora.name)
        adapter_weights.append(weight)

    print(f"Activating adapters: {adapter_names} with weights {adapter_weights}")
    pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)

    return pipe


def get_loras() -> list[dict]:
    return loras


def deactivate_loras(pipe):
    print("Unloading all LoRAs...")
    pipe.unload_lora_weights()
    return pipe


def calculate_optimal_dimensions(image):
    original_width, original_height = image.size
    MIN_ASPECT_RATIO = 9 / 16
    MAX_ASPECT_RATIO = 16 / 9
    FIXED_DIMENSION = 1024
    original_aspect_ratio = original_width / original_height
    if original_aspect_ratio > 1:
        width = FIXED_DIMENSION
        height = round(FIXED_DIMENSION / original_aspect_ratio)
    else:
        height = FIXED_DIMENSION
        width = round(FIXED_DIMENSION * original_aspect_ratio)
    width = (width // 8) * 8
    height = (height // 8) * 8
    calculated_aspect_ratio = width / height
    if calculated_aspect_ratio > MAX_ASPECT_RATIO:
        width = int((height * MAX_ASPECT_RATIO // 8) * 8)
    elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
        height = int((width / MIN_ASPECT_RATIO // 8) * 8)
    width = max(width, 576) if width == FIXED_DIMENSION else width
    height = max(height, 576) if height == FIXED_DIMENSION else height

    return width, height


@spaces.GPU(duration=45)
def inpaint(
        image,
        mask,
        prompt: str = "",
        seed: int = 0,
        num_inference_steps: int = 28,
        guidance_scale: int = 50,
        strength: float = 1.0
):
    image = image.convert("RGB")
    mask = mask.convert("L")
    width, height = calculate_optimal_dimensions(image)

    pipe.to("cuda")
    result = pipe(
        image=image,
        mask_image=mask,
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        strength=strength,
        generator=torch.Generator().manual_seed(seed)
    ).images[0]

    return result.convert("RGBA"), prompt, seed


def inpaint_api(
        image,
        mask,
        prompt: str,
        seed: int,
        num_inference_steps: int,
        guidance_scale: int,
        strength: float,
        flux_keywords: List[str] = None,
        loras_selected: List[Tuple[str, float]] = None
):
    flux_keywords = flux_keywords or []
    loras_selected = loras_selected or []

    # Convertir nombres a objetos LoRA
    selected_loras_with_weights = []

    for name, weight_value in loras_selected:
        try:
            # Convierte explícitamente el peso (que viene como string) a float
            weight = float(weight_value) 
        except (ValueError, TypeError):
            # Ignora si el valor no es un número válido (ej: None o string vacío)
            print(f"Valor de peso inválido '{weight_value}' para LoRA '{name}', omitiendo.")
            continue # Pasa al siguiente LoRA

        lora_obj = next((l for l in loras if l.display_name == name), None)
        
        # Ahora la comparación 'weight != 0.0' es segura (float con float)
        if lora_obj and weight != 0.0: 
            selected_loras_with_weights.append((lora_obj, weight))

    deactivate_loras(pipe)
    if selected_loras_with_weights:
        activate_loras(pipe, selected_loras_with_weights)

    # Construir prompt final
    final_prompt = ""

    if flux_keywords:
        final_prompt += ", ".join(flux_keywords) + ", "

    for lora, _ in selected_loras_with_weights:
        if lora.keyword:
            if isinstance(lora.keyword, str):
                final_prompt += lora.keyword + ", "
            else:
                final_prompt += ", ".join(lora.keyword) + ", "

    if final_prompt:
        final_prompt += "\n\n"
    final_prompt += prompt

    if not isinstance(seed, int) or seed < 0:
        seed = random.randint(0, MAX_SEED)

    return inpaint(
        image=image,
        mask=mask,
        prompt=final_prompt,
        seed=seed,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        strength=strength
    )


# ========================
# UI DIRECTA A inpaint_api
# ========================
with gr.Blocks(title="Flux.1 Fill dev Inpainting with LoRAs", theme=gr.themes.Soft()) as demo:
    gr.api(get_loras, api_name="get_loras")
    with gr.Row():
        
        with gr.Column(scale=2):
            prompt_input = gr.Text(
                label="Prompt", 
                lines=4, 
                value="a 25 years old woman"
            )
            
            seed_slider = gr.Slider(
                label="Seed", 
                minimum=-1, 
                maximum=MAX_SEED, 
                step=1,
                value=-1, 
                info="(-1 = Random)", 
                interactive=True
            )
            
            num_inference_steps_input = gr.Number(
                label="Inference steps", 
                value=40, interactive=True
            )
            
            guidance_scale_input = gr.Number(
                label="Guidance scale", 
                value=28, 
                interactive=True
            )
            
            strength_input = gr.Number(
                label="Strength", 
                value=1.0, 
                interactive=True, 
                maximum=1.0
            )

            gr.Markdown("### Flux Keywords")
            flux_keywords_input = gr.CheckboxGroup(
                choices=flux_keywords_available, 
                label="Flux Keywords"
            )

            if loras:
                gr.Markdown("### Available LoRAs")
                lora_names = [l.display_name for l in loras]
                loras_selected_input = gr.Dataframe(
                    type="array",
                    headers=["LoRA", "Weight"],
                    value=[[name, 0.0] for name in lora_names],
                    datatype=["str", "number"],  # Primera columna string, segunda número
                    interactive=[False, True],  # Solo la segunda columna editable
                    static_columns=[0],
                    label="LoRA selection (Weight 0 = disable)"
                )

        with gr.Column(scale=3):
            image_input = gr.Image(
                label="Image", 
                type="pil"
            )
            
            mask_input = gr.Image(
                label="Mask", 
                type="pil"
            )
            
            run_btn = gr.Button(
                "Run", 
                variant="primary"
            )

        with gr.Column(scale=3):
            result_image = gr.Image(
                label="Result"
            )
            
            used_prompt_box = gr.Text(
                label="Used prompt", 
                lines=4
            )
            
            used_seed_box = gr.Number(
                label="Used seed"
            )

    run_btn.click(
        fn=inpaint_api,
        inputs=[
            image_input,
            mask_input,
            prompt_input,
            seed_slider,
            num_inference_steps_input,
            guidance_scale_input,
            strength_input,
            flux_keywords_input,
            loras_selected_input
        ],
        outputs=[
            result_image, 
            used_prompt_box, 
            used_seed_box
        ],
        api_name="inpaint"
    )

if __name__ == "__main__":
    demo.launch(share=False, show_error=True)