import os 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 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") 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): if input_image is None: return None, 0 image, seed = infer( input_image=input_image, prompt=prompt, lora_adapter=lora_adapter, seed=0, randomize_seed=True, guidance_scale=1.0, steps=4 ) return image, seed 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-Studio**", elem_id="main-title") gr.Markdown("Perform diverse image edits using specialized [LoRA](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image-Edit-2509) adapters for the [FLUX.2-Klein-Model](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B) model.") with gr.Row(equal_height=True): with gr.Column(): input_image = gr.Image( label="Upload Image", type="pil", height=300, sources=["upload", "webcam", "clipboard"] ) with gr.Row(): prompt = gr.Text( label="Edit Prompt", max_lines=2, show_label=True, placeholder="e.g., a man with a red superhero mask" ) with gr.Row(): lora_adapter = gr.Dropdown( label="Choose Editing 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, visible=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="Output Image", interactive=False, format="png", height=420) used_seed = gr.Textbox(label="Used Seed", interactive=False, visible=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)