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": "examples/1.png",
"title": "Studio Ghibli",
"repo": "openfree/flux-chatgpt-ghibli-lora",
"trigger_word": "ghibli",
"weights": "pytorch_lora_weights.safetensors",
"likes": 0
},
{
"image": "examples/2.png",
"title": "Winslow Homer",
"repo": "openfree/winslow-homer",
"trigger_word": "homer",
"weights": "pytorch_lora_weights.safetensors",
"likes": 0
},
{
"image": "examples/3.png",
"title": "Van Gogh",
"repo": "openfree/van-gogh",
"trigger_word": "gogh",
"weights": "pytorch_lora_weights.safetensors",
"likes": 0
},
{
"image": "examples/4.png",
"title": "Paul Cรฉzanne",
"repo": "openfree/paul-cezanne",
"trigger_word": "Cezanne",
"weights": "pytorch_lora_weights.safetensors",
"likes": 0
},
{
"image": "examples/5.png",
"title": "Renoir",
"repo": "openfree/pierre-auguste-renoir",
"trigger_word": "Renoir",
"weights": "pytorch_lora_weights.safetensors",
"likes": 0
},
{
"image": "examples/6.png",
"title": "Claude Monet",
"repo": "openfree/claude-monet",
"trigger_word": "claude monet",
"weights": "pytorch_lora_weights.safetensors",
"likes": 0
},
{
"image": "examples/7.png",
"title": "Fantasy Art",
"repo": "openfree/myt-flux-fantasy",
"trigger_word": "fantasy",
"weights": "pytorch_lora_weights.safetensors",
"likes": 0
}
]
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 = """
/* Global background and container styling */
.gradio-container {
background: linear-gradient(135deg, #ffeef8 0%, #e6f3ff 25%, #fff4e6 50%, #f0e6ff 75%, #e6fff9 100%);
font-family: 'Inter', sans-serif;
}
/* Main app container */
#main_app {
display: flex;
gap: 24px;
padding: 20px;
background: rgba(255, 255, 255, 0.85);
backdrop-filter: blur(20px);
border-radius: 24px;
box-shadow: 0 10px 40px rgba(0, 0, 0, 0.08);
}
/* Box column styling */
#box_column {
min-width: 400px;
}
/* Gallery box with glassmorphism */
#gallery_box {
background: linear-gradient(135deg, rgba(255, 255, 255, 0.9) 0%, rgba(240, 248, 255, 0.9) 100%);
border-radius: 20px;
padding: 20px;
box-shadow: 0 8px 32px rgba(135, 206, 250, 0.2);
border: 1px solid rgba(255, 255, 255, 0.8);
}
/* Input image styling */
.image-container {
border-radius: 16px;
overflow: hidden;
box-shadow: 0 4px 20px rgba(0, 0, 0, 0.1);
}
/* Gallery styling */
#gallery {
overflow-y: scroll !important;
max-height: 400px;
padding: 12px;
background: rgba(255, 255, 255, 0.5);
border-radius: 16px;
scrollbar-width: thin;
scrollbar-color: #ddd6fe #f5f3ff;
}
#gallery::-webkit-scrollbar {
width: 8px;
}
#gallery::-webkit-scrollbar-track {
background: #f5f3ff;
border-radius: 10px;
}
#gallery::-webkit-scrollbar-thumb {
background: linear-gradient(180deg, #c7d2fe 0%, #ddd6fe 100%);
border-radius: 10px;
}
/* Selected LoRA text */
#selected_lora {
background: linear-gradient(135deg, #818cf8 0%, #a78bfa 100%);
-webkit-background-clip: text;
-webkit-text-fill-color: transparent;
background-clip: text;
font-weight: 700;
font-size: 18px;
text-align: center;
padding: 12px;
margin-bottom: 16px;
}
/* Prompt input field */
#prompt {
flex-grow: 1;
border: 2px solid transparent;
background: linear-gradient(white, white) padding-box,
linear-gradient(135deg, #a5b4fc 0%, #e9d5ff 100%) border-box;
border-radius: 12px;
padding: 12px 16px;
font-size: 16px;
transition: all 0.3s ease;
}
#prompt:focus {
box-shadow: 0 0 0 4px rgba(165, 180, 252, 0.25);
}
/* Run button with animated gradient */
#run_button {
background: linear-gradient(135deg, #a78bfa 0%, #818cf8 25%, #60a5fa 50%, #34d399 75%, #fbbf24 100%);
background-size: 200% 200%;
animation: gradient-shift 3s ease infinite;
color: white;
border: none;
padding: 12px 32px;
border-radius: 12px;
font-weight: 600;
font-size: 16px;
cursor: pointer;
transition: all 0.3s ease;
box-shadow: 0 4px 20px rgba(167, 139, 250, 0.4);
}
#run_button:hover {
transform: translateY(-2px);
box-shadow: 0 6px 30px rgba(167, 139, 250, 0.6);
}
@keyframes gradient-shift {
0% { background-position: 0% 50%; }
50% { background-position: 100% 50%; }
100% { background-position: 0% 50%; }
}
/* Custom LoRA card */
.custom_lora_card {
background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
border: 1px solid #fcd34d;
border-radius: 12px;
padding: 16px;
margin: 12px 0;
box-shadow: 0 4px 12px rgba(251, 191, 36, 0.2);
}
/* Result image container */
.output-image {
border-radius: 16px;
overflow: hidden;
box-shadow: 0 8px 32px rgba(0, 0, 0, 0.12);
margin-top: 20px;
}
/* Accordion styling */
.accordion {
background: rgba(249, 250, 251, 0.9);
border-radius: 12px;
border: 1px solid rgba(229, 231, 235, 0.8);
margin-top: 16px;
}
/* Slider styling */
.slider-container {
padding: 8px 0;
}
input[type="range"] {
background: linear-gradient(to right, #e0e7ff 0%, #c7d2fe 100%);
border-radius: 8px;
height: 6px;
}
/* Reuse button */
button:not(#run_button) {
background: linear-gradient(135deg, #f0abfc 0%, #c084fc 100%);
color: white;
border: none;
padding: 8px 20px;
border-radius: 8px;
font-weight: 500;
cursor: pointer;
transition: all 0.3s ease;
}
button:not(#run_button):hover {
transform: translateY(-1px);
box-shadow: 0 4px 16px rgba(192, 132, 252, 0.4);
}
/* Title styling */
h1 {
background: linear-gradient(135deg, #6366f1 0%, #a855f7 25%, #ec4899 50%, #f43f5e 75%, #f59e0b 100%);
-webkit-background-clip: text;
-webkit-text-fill-color: transparent;
background-clip: text;
text-align: center;
font-size: 3.5rem;
font-weight: 800;
margin-bottom: 8px;
text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.1);
}
h1 small {
display: block;
background: linear-gradient(135deg, #94a3b8 0%, #64748b 100%);
-webkit-background-clip: text;
-webkit-text-fill-color: transparent;
background-clip: text;
font-size: 1rem;
font-weight: 500;
margin-top: 8px;
}
/* Checkbox styling */
input[type="checkbox"] {
accent-color: #8b5cf6;
}
/* Label styling */
label {
color: #4b5563;
font-weight: 500;
}
/* Group containers */
.gr-group {
background: rgba(255, 255, 255, 0.7);
border-radius: 16px;
padding: 20px;
border: 1px solid rgba(255, 255, 255, 0.9);
box-shadow: 0 4px 16px rgba(0, 0, 0, 0.05);
}
"""
# Create Gradio interface
with gr.Blocks(css=css) as demo:
gr_flux_loras = gr.State(value=flux_loras_raw)
title = gr.HTML(
"""โจ Flux-Kontext FaceLORA
Transform your portraits with AI-powered style transfer ๐จ
""",
)
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 your portrait photo ๐ธ", type="pil", height=300)
gallery = gr.Gallery(
label="Choose Your Art Style",
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("Generate โจ", elem_id="run_button")
result = gr.Image(label="Your Artistic Portrait", 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]
)
# Initialize gallery
demo.load(
fn=classify_gallery,
inputs=[gr_flux_loras],
outputs=[gallery, gr_flux_loras]
)
demo.queue(default_concurrency_limit=None)
demo.launch()