app.py

import os
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import Flux2KleinPipeline
from huggingface_hub import login, hf_hub_download
from safetensors.torch import load_file
try:
    from color_matcher import ColorMatcher
except ImportError:
    import subprocess
    subprocess.check_call(["pip", "install", "color-matcher"])
    from color_matcher import ColorMatcher

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

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

# Login to HuggingFace for private repos if token is available
if "HF_TOKEN_LOGIN" in os.environ:
    login(token=os.environ["HF_TOKEN_LOGIN"])

# Model repository IDs
REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-4B"
REPO_ID_ALBEDO_LORA = "NightRaven109/kleinalbedo4B5ksteps"

# Load 4B Distilled model with Albedo LoRA
print("Loading 4B Distilled model...")
pipe = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)

print("Loading Albedo LoRA...")
# Download and load LoRA weights manually
lora_path = hf_hub_download(repo_id=REPO_ID_ALBEDO_LORA, filename="klein_albedo_5k_diffprompt.safetensors")
lora_state_dict = load_file(lora_path)

# Merge LoRA weights into model with proper key conversion
model_state_dict = pipe.transformer.state_dict()
lora_scale = 1.0
# Get hidden size from model weights
hidden_size = model_state_dict["transformer_blocks.0.attn.to_q.weight"].shape[0]
print(f"Model hidden size: {hidden_size}")

# Check LoRA dimensions
sample_qkv_lora_b = lora_state_dict.get("diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight")
if sample_qkv_lora_b is not None:
    lora_qkv_size = sample_qkv_lora_b.shape[0]
    lora_hidden_size = lora_qkv_size // 3
    print(f"LoRA QKV size: {lora_qkv_size}, LoRA hidden size: {lora_hidden_size}")

matched_keys = 0
unmatched_keys = []

# Store LoRA deltas for runtime adjustment
lora_deltas = {}
current_lora_scale = [lora_scale]  # Use list to allow modification in nested function

def apply_lora(model_state_dict, base_key, lora_a, lora_b, scale, store_delta=True):
    """Apply LoRA delta to model weights."""
    if base_key in model_state_dict:
        delta = (lora_b @ lora_a)
        if store_delta:
            lora_deltas[base_key] = delta.clone()
        model_state_dict[base_key] = model_state_dict[base_key] + (delta * scale).to(model_state_dict[base_key].dtype)
        return True
    return False

for key in lora_state_dict.keys():
    if not key.endswith(".lora_A.weight"):
        continue

    lora_b_key = key.replace(".lora_A.weight", ".lora_B.weight")
    if lora_b_key not in lora_state_dict:
        continue

    lora_a = lora_state_dict[key].to(dtype=dtype)
    lora_b = lora_state_dict[lora_b_key].to(dtype=dtype)

    # Remove prefix and get base name
    base = key.replace(".lora_A.weight", "")
    if base.startswith("diffusion_model."):
        base = base[len("diffusion_model."):]

    applied = False

    # Double blocks - image attention QKV (split into Q, K, V)
    if ".img_attn.qkv" in base:
        block_prefix = base.replace(".img_attn.qkv", "").replace("double_blocks.", "transformer_blocks.")
        # Split lora_b into Q, K, V parts
        q_b = lora_b[0:hidden_size, :]
        k_b = lora_b[hidden_size:2*hidden_size, :]
        v_b = lora_b[2*hidden_size:3*hidden_size, :]

        applied = apply_lora(model_state_dict, f"{block_prefix}.attn.to_q.weight", lora_a, q_b, lora_scale)
        applied |= apply_lora(model_state_dict, f"{block_prefix}.attn.to_k.weight", lora_a, k_b, lora_scale)
        applied |= apply_lora(model_state_dict, f"{block_prefix}.attn.to_v.weight", lora_a, v_b, lora_scale)

    # Double blocks - text attention QKV (split into add_q, add_k, add_v)
    elif ".txt_attn.qkv" in base:
        block_prefix = base.replace(".txt_attn.qkv", "").replace("double_blocks.", "transformer_blocks.")
        # Split lora_b into Q, K, V parts
        q_b = lora_b[0:hidden_size, :]
        k_b = lora_b[hidden_size:2*hidden_size, :]
        v_b = lora_b[2*hidden_size:3*hidden_size, :]

        applied = apply_lora(model_state_dict, f"{block_prefix}.attn.add_q_proj.weight", lora_a, q_b, lora_scale)
        applied |= apply_lora(model_state_dict, f"{block_prefix}.attn.add_k_proj.weight", lora_a, k_b, lora_scale)
        applied |= apply_lora(model_state_dict, f"{block_prefix}.attn.add_v_proj.weight", lora_a, v_b, lora_scale)

    # Double blocks - image attention output projection
    elif ".img_attn.proj" in base:
        new_key = base.replace("double_blocks.", "transformer_blocks.").replace(".img_attn.proj", ".attn.to_out.0") + ".weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)

    # Double blocks - text attention output projection
    elif ".txt_attn.proj" in base:
        new_key = base.replace("double_blocks.", "transformer_blocks.").replace(".txt_attn.proj", ".attn.to_add_out") + ".weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)

    # Single blocks - linear1 maps to to_qkv_mlp_proj
    elif "single_blocks." in base and ".linear1" in base:
        new_key = base.replace("single_blocks.", "single_transformer_blocks.").replace(".linear1", ".attn.to_qkv_mlp_proj") + ".weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)

    # Single blocks - linear2 maps to to_out
    elif "single_blocks." in base and ".linear2" in base:
        new_key = base.replace("single_blocks.", "single_transformer_blocks.").replace(".linear2", ".attn.to_out") + ".weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)

    # Modulation layers
    elif "double_stream_modulation_img.lin" in base:
        new_key = "double_stream_modulation_img.linear.weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)
    elif "double_stream_modulation_txt.lin" in base:
        new_key = "double_stream_modulation_txt.linear.weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)
    elif "single_stream_modulation.lin" in base:
        new_key = "single_stream_modulation.linear.weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)

    # Embedders
    elif "img_in" in base:
        new_key = "x_embedder.weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)
    elif "txt_in" in base:
        new_key = "context_embedder.weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)

    # Time embedding
    elif "time_in.in_layer" in base:
        new_key = "time_guidance_embed.timestep_embedder.linear_1.weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)
    elif "time_in.out_layer" in base:
        new_key = "time_guidance_embed.timestep_embedder.linear_2.weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)

    # Final layer
    elif "final_layer.linear" in base:
        new_key = "proj_out.weight"
        applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale)

    if applied:
        matched_keys += 1
    else:
        unmatched_keys.append(base)

print(f"LoRA keys matched: {matched_keys}")
print(f"LoRA keys unmatched: {len(unmatched_keys)}")
if unmatched_keys:
    print("Unmatched keys:", unmatched_keys)

pipe.transformer.load_state_dict(model_state_dict)
print("LoRA merged successfully!")
print(f"Stored {len(lora_deltas)} LoRA deltas for runtime adjustment")

def update_lora_scale(new_scale):
    """Update LoRA scale at runtime by adjusting weights."""
    global current_lora_scale
    old_scale = current_lora_scale[0]
    if new_scale == old_scale:
        return

    scale_diff = new_scale - old_scale
    state_dict = pipe.transformer.state_dict()

    for key, delta in lora_deltas.items():
        if key in state_dict:
            target_device = state_dict[key].device
            delta_on_device = delta.to(device=target_device, dtype=state_dict[key].dtype)
            state_dict[key] = state_dict[key] + (delta_on_device * scale_diff)

    pipe.transformer.load_state_dict(state_dict)
    current_lora_scale[0] = new_scale
    print(f"LoRA scale updated: {old_scale} -> {new_scale}")

pipe.to("cuda")


def update_dimensions_from_image(image_list):
    """Update width/height sliders based on uploaded image aspect ratio.
    Keeps one side at 1024 and scales the other proportionally, with both sides as multiples of 8."""
    if image_list is None or len(image_list) == 0:
        return 1024, 1024  # Default dimensions

    # Get the first image to determine dimensions
    img = image_list[0][0]  # Gallery returns list of tuples (image, caption)
    img_width, img_height = img.size

    aspect_ratio = img_width / img_height

    if aspect_ratio >= 1:  # Landscape or square
        new_width = 1024
        new_height = int(1024 / aspect_ratio)
    else:  # Portrait
        new_height = 1024
        new_width = int(1024 * aspect_ratio)

    # Round to nearest multiple of 8
    new_width = round(new_width / 8) * 8
    new_height = round(new_height / 8) * 8

    # Ensure within valid range (minimum 256, maximum 1024)
    new_width = max(256, min(1024, new_width))
    new_height = max(256, min(1024, new_height))

    return new_width, new_height


def reinhard_color_match(target, reference, strength=1.0):
    """Apply Reinhard color transfer from reference image to target image.
    Both inputs and output are PIL Images."""
    cm = ColorMatcher()
    target_np = np.array(target).astype(np.float32) / 255.0
    reference_np = np.array(reference).astype(np.float32) / 255.0
    matched_np = cm.transfer(src=target_np, ref=reference_np, method='reinhard')
    if strength < 1.0:
        matched_np = target_np + strength * (matched_np - target_np)
    matched_np = np.clip(matched_np * 255, 0, 255).astype(np.uint8)
    from PIL import Image
    return Image.fromarray(matched_np)


PROMPT = "Basecolor, no shadows, flat lighting, keep color"

@spaces.GPU(duration=85)
def infer(input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=1, guidance_scale=1.0, lora_weight=1.0, color_match_enabled=True, color_match_strength=1.0, progress=gr.Progress(track_tqdm=True)):

    # Update LoRA scale if changed
    update_lora_scale(lora_weight)

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Prepare image list (convert None or empty gallery to None)
    image_list = None
    if input_images is not None and len(input_images) > 0:
        image_list = []
        for item in input_images:
            image_list.append(item[0])

    progress(0.2, desc="Generating image...")

    generator = torch.Generator(device=device).manual_seed(seed)

    pipe_kwargs = {
        "prompt": PROMPT,
        "height": height,
        "width": width,
        "num_inference_steps": num_inference_steps,
        "guidance_scale": guidance_scale,
        "generator": generator,
    }
    
    # Add images if provided
    if image_list is not None:
        pipe_kwargs["image"] = image_list
    
    output_image = pipe(**pipe_kwargs).images[0]

    # Apply Reinhard color match to transfer input colors to output
    if color_match_enabled and image_list is not None and color_match_strength > 0:
        progress(0.8, desc="Applying Reinhard color match...")
        output_image = reinhard_color_match(output_image, image_list[0], strength=color_match_strength)

    # Create comparison tuple (input, output) for slider
    input_for_compare = image_list[0] if image_list else output_image
    comparison = (input_for_compare, output_image)

    return output_image, comparison, seed


css = """
#col-container {
    margin: 0 auto;
    max-width: 1200px;
}
.gallery-container img{
    object-fit: contain;
}
"""

with gr.Blocks() as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.2 [Klein] - 4B (Apache 2.0)
FLUX.2 [klein] is a fast, unified image generation and editing model designed for fast inference [[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B)], [[blog](https://bfl.ai/blog/flux-2)]
        """)
        with gr.Row():
            with gr.Column():
                run_button = gr.Button("Run", scale=1)

                input_images = gr.Gallery(
                    label="Input Image(s)",
                    type="pil",
                    columns=3,
                    rows=1,
                )

                with gr.Accordion("Advanced Settings", open=False):

                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )
                    
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                    
                    with gr.Row():
                        
                        width = gr.Slider(
                            label="Width",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=8,
                            value=1024,
                        )
                        
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=8,
                            value=1024,
                        )
                    
                    with gr.Row():
                        
                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=100,
                            step=1,
                            value=1,
                        )

                        guidance_scale = gr.Slider(
                            label="Guidance scale",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=1.0,
                        )

                    lora_weight = gr.Slider(
                        label="LoRA Weight",
                        minimum=0.0,
                        maximum=2.0,
                        step=0.05,
                        value=1.0,
                    )

                    color_match_enabled = gr.Checkbox(label="Reinhard Color Match", value=True)

                    color_match_strength = gr.Slider(
                        label="Color Match Strength",
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=1.0,
                    )


            with gr.Column():
                result = gr.Image(label="Result", show_label=False)
                comparison_slider = gr.ImageSlider(
                    label="Before / After Comparison",
                    type="pil",
                )

    # Auto-update dimensions when images are uploaded
    input_images.upload(
        fn=update_dimensions_from_image,
        inputs=[input_images],
        outputs=[width, height]
    )

    gr.on(
        triggers=[run_button.click],
        fn=infer,
        inputs=[input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, lora_weight, color_match_enabled, color_match_strength],
        outputs=[result, comparison_slider, seed]
    )

demo.launch(css=css)