import gradio as gr import numpy as np import spaces import torch import random import json import os from PIL import Image from diffusers import FluxKontextPipeline from diffusers.utils import load_image from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, list_repo_files from safetensors.torch import load_file import requests import re # Load Kontext model MAX_SEED = np.iinfo(np.int32).max pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda") # Load LoRA data flux_loras_raw = [ { "image": "https://huggingface.co/prithivMLmods/FLUX.1-Kontext-Cinematic-Relighting/resolve/main/images/1.png", "title": "Kontext Cinematic Relighting", "repo": "prithivMLmods/FLUX.1-Kontext-Cinematic-Relighting", "trigger_word": "Cinematic Relighting, Relight this portrait with warm, cinematic indoor lighting. Add soft amber highlights and gentle shadows to the face mimicking golden-hour light through a cozy room. Maintain natural skin texture and soft facial shadows, while enhancing eye catchlights for a vivid, lifelike look. Adjust white balance to a warmer tone, and slightly boost exposure to soften the darker midtones. Preserve the subject's pose and expression, and enhance the depth with gentle background bokeh and subtle filmic glow.", "weights": "FLUX.1-Kontext-Cinematic-Relighting.safetensors" }, ] print(f"Loaded {len(flux_loras_raw)} LoRAs") # Global variables for LoRA management current_lora = None lora_cache = {} def load_lora_weights(repo_id, weights_filename): """Load LoRA weights from HuggingFace""" try: # First try with the specified filename try: lora_path = hf_hub_download(repo_id=repo_id, filename=weights_filename) if repo_id not in lora_cache: lora_cache[repo_id] = lora_path return lora_path except Exception as e: print(f"Failed to load {weights_filename}, trying to find alternative LoRA files...") # If the specified file doesn't exist, try to find any .safetensors file from huggingface_hub import list_repo_files try: files = list_repo_files(repo_id) safetensors_files = [f for f in files if f.endswith(('.safetensors', '.bin')) and 'lora' in f.lower()] if not safetensors_files: # Try without 'lora' in filename safetensors_files = [f for f in files if f.endswith('.safetensors')] if safetensors_files: # Try the first available file for file in safetensors_files: try: print(f"Trying alternative file: {file}") lora_path = hf_hub_download(repo_id=repo_id, filename=file) if repo_id not in lora_cache: lora_cache[repo_id] = lora_path print(f"Successfully loaded alternative LoRA file: {file}") return lora_path except: continue print(f"No suitable LoRA files found in {repo_id}") return None except Exception as list_error: print(f"Error listing files in repo {repo_id}: {list_error}") return None except Exception as e: print(f"Error loading LoRA from {repo_id}: {e}") return None def update_selection(selected_state: gr.SelectData, flux_loras): """Update UI when a LoRA is selected""" if selected_state.index >= len(flux_loras): return "### No LoRA selected", gr.update(), None lora = flux_loras[selected_state.index] lora_title = lora["title"] lora_repo = lora["repo"] trigger_word = lora["trigger_word"] # Create a more informative selected text updated_text = f"### 🎨 Selected Style: {lora_title}" new_placeholder = f"Describe additional details, e.g., 'wearing a red hat' or 'smiling'" return updated_text, gr.update(placeholder=new_placeholder), selected_state.index def get_huggingface_lora(link): """Download LoRA from HuggingFace link""" split_link = link.split("/") if len(split_link) == 2: try: model_card = ModelCard.load(link) trigger_word = model_card.data.get("instance_prompt", "") # Try to find the correct safetensors file files = list_repo_files(link) safetensors_files = [f for f in files if f.endswith('.safetensors')] # Prioritize files with 'lora' in the name lora_files = [f for f in safetensors_files if 'lora' in f.lower()] if lora_files: safetensors_file = lora_files[0] elif safetensors_files: safetensors_file = safetensors_files[0] else: # Try .bin files as fallback bin_files = [f for f in files if f.endswith('.bin') and 'lora' in f.lower()] if bin_files: safetensors_file = bin_files[0] else: safetensors_file = "pytorch_lora_weights.safetensors" # Default fallback print(f"Found LoRA file: {safetensors_file} in {link}") return split_link[1], safetensors_file, trigger_word except Exception as e: print(f"Error in get_huggingface_lora: {e}") # Try basic detection try: files = list_repo_files(link) safetensors_file = next((f for f in files if f.endswith('.safetensors')), "pytorch_lora_weights.safetensors") return split_link[1], safetensors_file, "" except: raise Exception(f"Error loading LoRA: {e}") else: raise Exception("Invalid HuggingFace repository format") def load_custom_lora(link): """Load custom LoRA from user input""" if not link: return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### 🎨 Select an art style from the gallery", None try: repo_name, weights_file, trigger_word = get_huggingface_lora(link) card = f'''
βœ… Custom LoRA Loaded!

{repo_name}

{"Trigger: "+trigger_word+"" if trigger_word else "No trigger word found"}
''' custom_lora_data = { "repo": link, "weights": weights_file, "trigger_word": trigger_word } return gr.update(visible=True), card, gr.update(visible=True), custom_lora_data, gr.Gallery(selected_index=None), f"🎨 Custom Style: {repo_name}", None except Exception as e: return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### 🎨 Select an art style from the gallery", None def remove_custom_lora(): """Remove custom LoRA""" return "", gr.update(visible=False), gr.update(visible=False), None, None def classify_gallery(flux_loras): """Sort gallery by likes""" try: sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True) gallery_items = [] for item in sorted_gallery: if "image" in item and "title" in item: image_path = item["image"] title = item["title"] # Simply use the path as-is for Gradio to handle gallery_items.append((image_path, title)) print(f"Added to gallery: {image_path} - {title}") print(f"Total gallery items: {len(gallery_items)}") return gallery_items, sorted_gallery except Exception as e: print(f"Error in classify_gallery: {e}") import traceback traceback.print_exc() return [], [] def infer_with_lora_wrapper(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)): """Wrapper function to handle state serialization""" return infer_with_lora(input_image, prompt, selected_index, custom_lora, seed, randomize_seed, guidance_scale, lora_scale, flux_loras, progress) @spaces.GPU def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)): """Generate image with selected LoRA""" global current_lora, pipe # Check if input image is provided if input_image is None: gr.Warning("Please upload your portrait photo first! πŸ“Έ") return None, seed, gr.update(visible=False) if randomize_seed: seed = random.randint(0, MAX_SEED) # Determine which LoRA to use lora_to_use = None if custom_lora: lora_to_use = custom_lora elif selected_index is not None and flux_loras and selected_index < len(flux_loras): lora_to_use = flux_loras[selected_index] # Load LoRA if needed if lora_to_use and lora_to_use != current_lora: try: # Unload current LoRA if current_lora: pipe.unload_lora_weights() print(f"Unloaded previous LoRA") # Load new LoRA repo_id = lora_to_use.get("repo", "unknown") weights_file = lora_to_use.get("weights", "pytorch_lora_weights.safetensors") print(f"Loading LoRA: {repo_id} with weights: {weights_file}") lora_path = load_lora_weights(repo_id, weights_file) if lora_path: pipe.load_lora_weights(lora_path, adapter_name="selected_lora") pipe.set_adapters(["selected_lora"], adapter_weights=[lora_scale]) print(f"Successfully loaded: {lora_path} with scale {lora_scale}") current_lora = lora_to_use else: print(f"Failed to load LoRA from {repo_id}") gr.Warning(f"Failed to load {lora_to_use.get('title', 'style')}. Please try a different art style.") return None, seed, gr.update(visible=False) except Exception as e: print(f"Error loading LoRA: {e}") # Continue without LoRA else: if lora_to_use: print(f"Using already loaded LoRA: {lora_to_use.get('repo', 'unknown')}") try: # Convert image to RGB input_image = input_image.convert("RGB") except Exception as e: print(f"Error processing image: {e}") gr.Warning("Error processing the uploaded image. Please try a different photo. πŸ“Έ") return None, seed, gr.update(visible=False) # Check if LoRA is selected if lora_to_use is None: gr.Warning("Please select an art style from the gallery first! 🎨") return None, seed, gr.update(visible=False) # Add trigger word to prompt trigger_word = lora_to_use.get("trigger_word", "") # Special handling for different art styles if trigger_word == "ghibli": prompt = f"Create a Studio Ghibli anime style portrait of the person in the photo, {prompt}. Maintain the facial identity while transforming into whimsical anime art style." elif trigger_word == "homer": prompt = f"Paint the person in Winslow Homer's American realist style, {prompt}. Keep facial features while applying watercolor and marine art techniques." elif trigger_word == "gogh": prompt = f"Transform the portrait into Van Gogh's post-impressionist style with swirling brushstrokes, {prompt}. Maintain facial identity with expressive colors." elif trigger_word == "Cezanne": prompt = f"Render the person in Paul CΓ©zanne's geometric post-impressionist style, {prompt}. Keep facial structure while applying structured brushwork." elif trigger_word == "Renoir": prompt = f"Paint the portrait in Pierre-Auguste Renoir's impressionist style with soft light, {prompt}. Maintain identity with luminous skin tones." elif trigger_word == "claude monet": prompt = f"Create an impressionist portrait in Claude Monet's style with visible brushstrokes, {prompt}. Keep facial features while using light and color." elif trigger_word == "fantasy": prompt = f"Transform into an epic fantasy character portrait, {prompt}. Maintain facial identity while adding magical and fantastical elements." elif trigger_word == ", How2Draw": prompt = f"create a How2Draw sketch of the person of the photo {prompt}, maintain the facial identity of the person and general features" elif trigger_word == ", video game screenshot in the style of THSMS": prompt = f"create a video game screenshot in the style of THSMS with the person from the photo, {prompt}. maintain the facial identity of the person and general features" else: prompt = f"convert the style of this portrait photo to {trigger_word} while maintaining the identity of the person. {prompt}. Make sure to maintain the person's facial identity and features, while still changing the overall style to {trigger_word}." try: image = pipe( image=input_image, prompt=prompt, guidance_scale=guidance_scale, generator=torch.Generator().manual_seed(seed), ).images[0] return image, seed, gr.update(visible=True) except Exception as e: print(f"Error during inference: {e}") return None, seed, gr.update(visible=False) # CSS styling with beautiful gradient pastel design css = ''' #gen_btn{height: 100%} #gen_column{align-self: stretch} #title{text-align: center} #title h1{font-size: 3em; display:inline-flex; align-items:center} #title img{width: 100px; margin-right: 0.5em} #gallery .grid-wrap{height: 10vh} #lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%} .card_internal{display: flex;height: 100px;margin-top: .5em} .card_internal img{margin-right: 1em} .styler{--form-gap-width: 0px !important} #progress{height:30px} #progress .generating{display:none} .progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px} .progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out} ''' # Create Gradio interface with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo: gr_flux_loras = gr.State(value=flux_loras_raw) title = gr.HTML( """

Flux Kontext DLC 🎈

""", ) selected_state = gr.State(value=None) custom_loaded_lora = gr.State(value=None) with gr.Row(elem_id="main_app"): with gr.Column(scale=4, elem_id="box_column"): with gr.Group(elem_id="gallery_box"): input_image = gr.Image(label="Upload an image for editing", type="pil", height=260) gallery = gr.Gallery( label="Choose the Flux Kontext LoRA", allow_preview=False, columns=3, elem_id="gallery", show_share_button=False, height=400 ) custom_model = gr.Textbox( label="πŸ”— Or use a custom LoRA from HuggingFace", placeholder="e.g., username/lora-name", visible=True ) custom_model_card = gr.HTML(visible=False) custom_model_button = gr.Button("Remove custom LoRA", visible=False) with gr.Column(scale=5): with gr.Row(): prompt = gr.Textbox( label="Additional Details (optional)", show_label=False, lines=1, max_lines=1, placeholder="Describe additional details, e.g., 'wearing a red hat' or 'smiling'", elem_id="prompt" ) run_button = gr.Button("Edit Image", elem_id="run_button") result = gr.Image(label="Your Kontext Edited Image", interactive=False) reuse_button = gr.Button("Reuse this image", visible=False) with gr.Accordion("Advanced Settings", open=False): lora_scale = gr.Slider( label="Style Strength", minimum=0, maximum=2, step=0.1, value=1.0, info="How strongly to apply the art style (1.0 = balanced)" ) seed = gr.Slider( label="Random Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, info="Set to 0 for random results" ) randomize_seed = gr.Checkbox(label="Randomize seed for each generation", value=True) guidance_scale = gr.Slider( label="Image Guidance", minimum=1, maximum=10, step=0.1, value=2.5, info="How closely to follow the input image (lower = more creative)" ) prompt_title = gr.Markdown( value="### Select an art style from the gallery", visible=True, elem_id="selected_lora", ) # Event handlers custom_model.input( fn=load_custom_lora, inputs=[custom_model], outputs=[custom_model_card, custom_model_card, custom_model_button, custom_loaded_lora, gallery, prompt_title, selected_state], ) custom_model_button.click( fn=remove_custom_lora, outputs=[custom_model, custom_model_button, custom_model_card, custom_loaded_lora, selected_state] ) gallery.select( fn=update_selection, inputs=[gr_flux_loras], outputs=[prompt_title, prompt, selected_state], show_progress=False ) gr.on( triggers=[run_button.click, prompt.submit], fn=infer_with_lora_wrapper, inputs=[input_image, prompt, selected_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, gr_flux_loras], outputs=[result, seed, reuse_button] ) reuse_button.click( fn=lambda image: image, inputs=[result], outputs=[input_image] ) demo.load( fn=classify_gallery, inputs=[gr_flux_loras], outputs=[gallery, gr_flux_loras] ) demo.queue(default_concurrency_limit=None) demo.launch()