# frontend.py
import gradio as gr
import sys
import os
import spaces
# Add the parent directory to sys.path
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, parent_dir)
#print(sys.path) #DEBUG
from flux_app.backend import ModelManager # Absolute import
from flux_app.config import MAX_SEED # Absolute import
from flux_app.lora_handling import (
add_custom_lora, remove_custom_lora, prepare_prompt,
unload_lora_weights, load_lora_weights_into_pipeline, update_selection
)
from flux_app.utilities import randomize_seed_if_needed, calculateDuration # Absolute import
# Import the prompt enhancer function
from flux_app.enhance import generate as enhance_generate
# Dummy loras data for initial UI setup.
initial_loras = [
{"image": "placeholder.jpg", "title": "Placeholder LoRA", "repo": "placeholder/repo", "weights": None, "trigger_word": ""},
]
class Frontend:
def __init__(self, model_manager: ModelManager):
self.model_manager = model_manager
self.loras = initial_loras
self.load_initial_loras()
self.css = self.define_css()
def define_css(self):
# A cleaner, professional CSS styling.
return '''
/* Title Styling */
#title {
text-align: center;
margin-bottom: 20px;
}
#title h1 {
font-size: 2.5rem;
margin: 0;
color: #333;
}
/* Button and Column Styling */
#gen_btn {
width: 100%;
padding: 12px;
font-weight: bold;
border-radius: 5px;
}
#gen_column {
display: flex;
align-items: center;
justify-content: center;
}
/* Gallery and List Styling */
#gallery .grid-wrap {
margin-top: 15px;
}
#lora_list {
background-color: #f5f5f5;
padding: 10px;
border-radius: 4px;
font-size: 0.9rem;
}
.card_internal {
display: flex;
align-items: center;
height: 100px;
margin-top: 10px;
}
.card_internal img {
margin-right: 10px;
}
.styler {
--form-gap-width: 0px !important;
}
/* Progress Bar Styling */
.progress-container {
width: 100%;
height: 20px;
background-color: #e0e0e0;
border-radius: 10px;
overflow: hidden;
margin-bottom: 20px;
}
.progress-bar {
height: 100%;
background-color: #4f46e5;
transition: width 0.3s ease-in-out;
width: calc(var(--current) / var(--total) * 100%);
}
'''
def load_initial_loras(self):
try:
from flux_app.lora import loras as loras_list # Absolute import
self.loras = loras_list
except ImportError:
print("Warning: lora.py not found, using placeholder LoRAs.")
pass
@spaces.GPU(duration=100)
def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index,
randomize_seed, seed, width, height, lora_scale, use_enhancer,
progress=gr.Progress(track_tqdm=True)):
seed = randomize_seed_if_needed(randomize_seed, seed, MAX_SEED)
# Prepare the initial prompt (using LoRA info if needed)
prompt_mash = prepare_prompt(prompt, selected_index, self.loras)
enhanced_text = ""
# If prompt enhancer is enabled, first run it to improve the prompt.
if use_enhancer:
# Stream the enhanced prompt (this will update the enhanced prompt textbox)
for enhanced_chunk in enhance_generate(prompt_mash):
enhanced_text = enhanced_chunk
# Yield an update with no image yet and the current enhanced prompt.
yield None, seed, gr.update(visible=False), enhanced_text
# Use the final enhanced prompt as the prompt for image generation.
prompt_mash = enhanced_text
else:
# Ensure the enhanced prompt textbox remains cleared.
enhanced_text = ""
# Continue with the image generation process.
selected_lora = self.loras[selected_index]
unload_lora_weights(self.model_manager.pipe, self.model_manager.pipe_i2i)
pipe_to_use = self.model_manager.pipe_i2i if image_input is not None else self.model_manager.pipe
load_lora_weights_into_pipeline(pipe_to_use, selected_lora["repo"], selected_lora.get("weights"))
if image_input is not None:
final_image = self.model_manager.generate_image_to_image(
prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed
)
yield final_image, seed, gr.update(visible=False), enhanced_text
else:
image_generator = self.model_manager.generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale)
final_image = None
step_counter = 0
for image in image_generator:
step_counter += 1
final_image = image
progress_bar = f''
yield image, seed, gr.update(value=progress_bar, visible=True), enhanced_text
yield final_image, seed, gr.update(value=progress_bar, visible=False), enhanced_text
def create_ui(self):
# Using a base theme for a clean and professional look.
with gr.Blocks(theme=gr.themes.Base(), css=self.css, title="Flux LoRA Generation") as app:
title = gr.HTML(
"""Flux LoRA Generation
""",
elem_id="title",
)
selected_index = gr.State(None)
with gr.Row():
with gr.Column(scale=3):
prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Choose the LoRA and type the prompt")
with gr.Column(scale=1, elem_id="gen_column"):
generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
with gr.Row():
with gr.Column():
selected_info = gr.Markdown("")
gallery = gr.Gallery(
[(item["image"], item["title"]) for item in self.loras],
label="LoRA Collection",
allow_preview=False,
columns=3,
elem_id="gallery",
show_share_button=False
)
with gr.Group():
custom_lora = gr.Textbox(label="Enter Custom LoRA", placeholder="prithivMLmods/Canopus-LoRA-Flux-Anime")
gr.Markdown("[Check the list of FLUX LoRA's](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
custom_lora_info = gr.HTML(visible=False)
custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
with gr.Column():
progress_bar = gr.Markdown(elem_id="progress", visible=False)
result = gr.Image(label="Generated Image")
with gr.Row():
with gr.Accordion("Advanced Settings", open=False):
with gr.Row():
input_image = gr.Image(label="Input image", type="filepath")
image_strength = gr.Slider(label="Denoise Strength", info="Lower means more image influence", minimum=0.1, maximum=1.0, step=0.01, value=0.75)
with gr.Column():
with gr.Row():
cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
with gr.Row():
width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
with gr.Row():
randomize_seed = gr.Checkbox(True, label="Randomize seed")
seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
with gr.Row():
use_enhancer = gr.Checkbox(value=False, label="Use Prompt Enhancer")
show_enhanced_prompt = gr.Checkbox(value=False, label="Display Enhanced Prompt")
# Enhanced prompt textbox (hidden by default)
enhanced_prompt_box = gr.Textbox(label="Enhanced Prompt", visible=False)
gallery.select(
update_selection,
inputs=[width, height, gr.State(self.loras)],
outputs=[prompt, selected_info, selected_index, width, height]
)
custom_lora.input(
add_custom_lora,
inputs=[custom_lora, gr.State(self.loras)],
outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
)
custom_lora_button.click(
remove_custom_lora,
outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
)
# Toggle the visibility of the enhanced prompt textbox based on the checkbox state.
show_enhanced_prompt.change(fn=lambda show: gr.update(visible=show),
inputs=show_enhanced_prompt,
outputs=enhanced_prompt_box)
gr.on(
triggers=[generate_button.click, prompt.submit],
fn=self.run_lora,
inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index,
randomize_seed, seed, width, height, lora_scale, use_enhancer],
outputs=[result, seed, progress_bar, enhanced_prompt_box]
)
# Credits section added at the bottom
with gr.Row():
gr.HTML("")
return app
if __name__ == "__main__":
model_manager = ModelManager()
frontend = Frontend(model_manager)
app = frontend.create_ui()
app.queue()
app.launch()