app.py

import os
import gc
import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image
from typing import Iterable

from diffusers import Flux2KleinPipeline
from diffusers.utils import load_image
from huggingface_hub import hf_hub_download

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

from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

colors.orange_red = colors.Color(
    name="orange_red", c50="#FFF0E5", c100="#FFE0CC", c200="#FFC299", c300="#FFA366",
    c400="#FF8533", c500="#FF4500", c600="#E63E00", c700="#CC3700", c800="#B33000",
    c900="#992900", c950="#802200",
)

class OrangeRedTheme(Soft):
    def __init__(
        self, *, primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.orange_red,
        neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue,
            text_size=text_size, font=font, font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
            slider_color="*secondary_500",
            slider_color_dark="*secondary_600",
            block_title_text_weight="600", block_border_width="3px",
            block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px", color_accent_soft="*primary_100",
            block_label_background_fill="*primary_200",
        )

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

ADAPTER_SPECS = {
    "Simple-Tuner": {
        "repo": "markury/flux2k9b-simpletuner-lora-loona",
        "weights": "pytorch_lora_weights.safetensors",
        "adapter_name": "simple-tuner"
    },
    "Klein-Delight-Style": {
        "repo": "linoyts/Flux2-Klein-Delight-LoRA",
        "weights": "pytorch_lora_weights.safetensors",
        "adapter_name": "klein-delight"
    }
}

LOADED_ADAPTERS = set()

print("Loading FLUX.2 Klein 9B model base...")
pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B",
    torch_dtype=torch.bfloat16
).to(device)
print("Base Model loaded successfully.")

def update_dimensions_on_upload(image):
    """Resizes image to be divisible by 16 to avoid tensor mismatch errors in FLUX."""
    if image is None:
        return 1024, 1024
    
    original_width, original_height = image.size
    
    # Simple logic to keep aspect ratio but snap to multiples of 16
    # You can adjust the max size (e.g., 1024) as needed
    scale = min(1024 / original_width, 1024 / original_height)
    new_width = int(original_width * scale)
    new_height = int(original_height * scale)
    
    new_width = (new_width // 16) * 16
    new_height = (new_height // 16) * 16
    
    return new_width, new_height

@spaces.GPU
def infer(
    input_image, 
    prompt, 
    lora_adapter, 
    seed=42, 
    randomize_seed=True, 
    guidance_scale=1.0, 
    steps=4, 
    progress=gr.Progress(track_tqdm=True)
):
    gc.collect()
    torch.cuda.empty_cache()

    if not input_image:
        raise gr.Error("Please upload an image to apply a style to.")

    spec = ADAPTER_SPECS.get(lora_adapter)
    if spec:
        adapter_name = spec["adapter_name"]
        
        if adapter_name not in LOADED_ADAPTERS:
            print(f"--- Downloading and Loading Adapter: {lora_adapter} ---")
            try:
                pipe.load_lora_weights(
                    spec["repo"], 
                    weight_name=spec["weights"], 
                    adapter_name=adapter_name
                )
                LOADED_ADAPTERS.add(adapter_name)
            except Exception as e:
                raise gr.Error(f"Failed to load adapter {lora_adapter}: {e}")
        else:
            print(f"--- Adapter {lora_adapter} is already loaded. ---")
            
        print(f"Activating LoRA: {adapter_name}")
        pipe.set_adapters([adapter_name], adapter_weights=[1.0])
    else:
        print("No valid LoRA selected or found. Disabling adapters.")
        pipe.disable_lora()

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    width, height = update_dimensions_on_upload(input_image)
    processed_input = input_image.resize((width, height), Image.LANCZOS).convert("RGB")
    
    # 4. Inference
    try:
        image = pipe(
            image=processed_input, 
            prompt=prompt,
            guidance_scale=guidance_scale,
            width=width,
            height=height,
            num_inference_steps=steps,
            generator=torch.Generator(device=device).manual_seed(seed),
        ).images[0]
        
        return image, seed

    except Exception as e:
        raise gr.Error(f"Inference failed: {e}")
    finally:
        gc.collect()
        torch.cuda.empty_cache()

@spaces.GPU
def infer_example(input_image, prompt, lora_adapter):
    # Use a fixed seed for reproducible examples
    if input_image is None: 
        return None, 0
        
    image, seed = infer(
        input_image=input_image, 
        prompt=prompt, 
        lora_adapter=lora_adapter, 
        seed=12345, 
        randomize_seed=False,
        guidance_scale=4.0,
        steps=4
    )
    return image, seed

# --- UI Layout ---
css="""
#col-container { margin: 0 auto; max-width: 960px; }
#main-title h1 { font-size: 2.2em !important; }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **FLUX.2 Klein LoRA Stylizer**", elem_id="main-title")
        gr.Markdown(
            "Apply creative styles to your images using **FLUX.2-klein-9B** and specialized LoRA adapters. "
            "Adapters are loaded lazily (on-demand) to save resources."
        )
        
        with gr.Row(equal_height=True):
            with gr.Column():
                input_image = gr.Image(
                    label="Upload Image", 
                    type="pil", 
                    height=290, 
                    sources=["upload", "webcam", "clipboard"]
                )
                prompt = gr.Text(
                    label="Guiding Prompt", 
                    show_label=True, 
                    placeholder="e.g., a man with a red superhero mask"
                )
                
                lora_adapter = gr.Dropdown(
                    label="Choose a Creative Style",
                    choices=list(ADAPTER_SPECS.keys()),
                    value="Klein-Delight-Style"
                )

                run_button = gr.Button("Apply Style", variant="primary")

                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)
                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0)       
                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
                    
            with gr.Column():
                output_image = gr.Image(label="Stylized Image", interactive=False, format="png", height=450)
                used_seed = gr.Textbox(label="Used Seed", interactive=False)

        gr.Examples(
            examples=[
                ["examples/animal.jpg", "a cute red panda, charming and delightful illustration, soft lighting", "Klein-Delight-Style"],
                ["examples/portrait.jpg", "cinematic lighting, high contrast", "Simple-Tuner"],
            ],
            inputs=[input_image, prompt, lora_adapter],
            outputs=[output_image, used_seed],
            fn=infer_example,
            cache_examples=False,
        )
            
    run_button.click(
        fn=infer,
        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
        outputs=[output_image, used_seed]
    )

if __name__ == "__main__":
    demo.queue().launch(css=css, theme=orange_red_theme, show_error=True)