| 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) |
