flux_app/frontend.py

# 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=300)
    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'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
                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(
                """<h1>Flux LoRA Generation</h1>""",
                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("<div style='text-align:center; font-size:0.9em; margin-top:20px;'>Credits: <a href='https://ruslanmv.com' target='_blank'>ruslanmv.com</a></div>")
            
            return app

if __name__ == "__main__":
    model_manager = ModelManager()
    frontend = Frontend(model_manager)
    app = frontend.create_ui()
    app.queue()
    app.launch()