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)