gradio_app.py

import os
import zipfile
from pathlib import Path

import gradio as gr

from database import (
    get_user_credits,
    update_user_credits,
    get_lora_models_info,
    get_user_lora_models
)

from services.image_generation import generate_image
from services.train_lora import lora_pipeline
from utils.image_utils import url_to_pil_image
from utils.file_utils import load_file_content

LORA_MODELS = get_lora_models_info()
if not isinstance(LORA_MODELS, list):
    raise ValueError("Expected loras_models to be a list of dictionaries.")

BASE_DIR = Path(__file__).parent
LOGIN_CSS_PATH = BASE_DIR / 'static/css/login.css'
MAIN_CSS_PATH = BASE_DIR / 'static/css/main.css'
LANDING_HTML_PATH = BASE_DIR / 'static/html/landing.html'
MAIN_HEADER_PATH = BASE_DIR / 'static/html/main_header.html'

LOGIN_CSS = load_file_content(LOGIN_CSS_PATH)
MAIN_CSS = load_file_content(MAIN_CSS_PATH)
LANDING_PAGE = load_file_content(LANDING_HTML_PATH)
MAIN_HEADER = load_file_content(MAIN_HEADER_PATH)

def load_user_models(request: gr.Request):
    user = request.session.get('user')
    print(user)
    if user:
        user_models = get_user_lora_models(user['id'])
        if user_models:
            return [(item.get("image_url", "assets/logo.jpg"), item["lora_name"]) for item in user_models]
    return []

def update_selection(evt: gr.SelectData, gallery_type: str, width, height):
    if gallery_type == "user":
        selected_lora = {"lora_name": "custom", "trigger_word": "custom"}
    else:
        selected_lora = LORA_MODELS[evt.index]
    new_placeholder = f"Enter a prompt for {selected_lora['lora_name']}"
    trigger_word = selected_lora["trigger_word"]
    updated_text = f"#### Trigger Word: {trigger_word} ✨"
    
    if "aspect" in selected_lora:
        if selected_lora["aspect"] == "portrait":
            width, height = 768, 1024
        elif selected_lora["aspect"] == "landscape":
            width, height = 1024, 768

    return gr.update(placeholder=new_placeholder), updated_text, evt.index, width, height, gallery_type

def compress_and_train(request: gr.Request, files, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate):
    if not files:
        return "No Images. Please, upload some images to start training"
    
    user = request.session.get('user')
    
    _, training_credits = get_user_credits(user['id'])
    
    if training_credits <= 0:
        raise gr.Error("You ran out of credtis. Please buy more to continue")
    
    if not user:
        raise gr.Error("User not authenticated. Please log in.")

    user_id = user['id']
    # Create a directory in the user's home folder
    output_dir = os.path.expanduser("~/gradio_training_data")
    os.makedirs(output_dir, exist_ok=True)

    # Create a zip file in the output directory
    zip_path = os.path.join(output_dir, "training_data.zip")
    
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file_info in files:
            file_path = file_info[0]  # The first element of the tuple is the file path
            file_name = os.path.basename(file_path)
            zipf.write(file_path, file_name)
    
    print(f"Zip file created at: {zip_path}")
        
    print(f'[INFO] Procesando {trigger_word}')
    # Now call the train_lora function with the zip file path
    result = lora_pipeline(user_id,
                            zip_path, 
                            model_name, 
                            trigger_word=trigger_word, 
                            steps=train_steps, 
                            lora_rank=lora_rank, 
                            batch_size=batch_size, 
                            autocaption=True, 
                            learning_rate=learning_rate)
    
    new_training_credits = training_credits - 1
    update_user_credits(user['id'], user['generation_credits'], new_training_credits)
    
    # Update session data
    user['training_credits'] = new_training_credits
    request.session['user'] = user
    
    return gr.Info("Your model is training. In about 20 minutes, it will be ready for you to test in 'Generation"), new_training_credits
            
def run_lora(request: gr.Request, prompt, cfg_scale, steps, selected_index, selected_gallery, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
    user = request.session.get('user')
    if not user:
        raise gr.Error("User not authenticated. Please log in.")
    lora_models = get_user_lora_models(user['id'])
    print(f'Selected gallery: {selected_gallery}')
    if selected_gallery == "user":
        lora_models = get_user_lora_models(user['id'])
        print('Using user models')
    else:  # public
        lora_models = get_lora_models_info()
        print('Using public models')
    print(f'Selected index: {selected_index}')
    if selected_index is None:
        selected_lora = None
    else:
        selected_lora = lora_models[selected_index]

    generation_credits, _ = get_user_credits(user['id'])
    if selected_lora:
        print(f"Selected Lora: {selected_lora['lora_name']}")
        model_name = selected_lora['lora_name']
        use_default = False
    else:
        model_name = "black-forest-labs/flux-pro"
        print(f"Using default Lora: {model_name}")
        use_default = True
    if generation_credits <= 0:
        raise gr.Error("Ya no tienes creditos disponibles. Compra para continuar.")
    
    image_url = generate_image(model_name, prompt, steps, cfg_scale, width, height, lora_scale, progress, use_default)
    image = url_to_pil_image(image_url)
    
    # Update user's credits
    new_generation_credits = generation_credits - 1
    update_user_credits(user['id'], new_generation_credits, user['train_credits'])
    
    # Update session data
    user['generation_credits'] = new_generation_credits
    request.session['user'] = user
    
    print(f"Generation credits remaining: {new_generation_credits}")
    
    return image, new_generation_credits

def display_credits(request: gr.Request):
    user = request.session.get('user')
    if user:
        generation_credits, train_credits = get_user_credits(user['id'])
        return generation_credits, train_credits
    return 0, 0

def load_greet_and_credits(request: gr.Request):
    greeting = greet(request)
    generation_credits, train_credits = display_credits(request)
    return greeting, generation_credits, train_credits
    
def greet(request: gr.Request):
    user = request.session.get('user')
    if user:
        with gr.Column():  
            with gr.Row():
                greeting = f"Hola 👋 {user['given_name']}!"
            return f"{greeting}\n"
    return "OBTU AI. Please log in."

with gr.Blocks(theme=gr.themes.Soft(), css=LOGIN_CSS) as login_demo:
    with gr.Column(elem_id="google-btn-container", elem_classes="google-btn-container svelte-vt1mxs gap"):
        btn = gr.Button("Sign In with Google", elem_classes="login-with-google-btn")
    _js_redirect = """
    () => {
        url = '/login' + window.location.search;
        window.open(url, '_blank');
    }
    """
    btn.click(None, js=_js_redirect)
    gr.HTML(LANDING_PAGE)
        

header = '<script src="https://cdn.lordicon.com/lordicon.js"></script>'

with gr.Blocks(theme=gr.themes.Soft(), head=header, css=MAIN_CSS) as main_demo:
    title = gr.HTML(MAIN_HEADER)
    
    with gr.Column(elem_id="logout-btn-container"):
        gr.Button("Logout", link="/logout", elem_id="logout_btn")

    
    greetings = gr.Markdown("Loading user information...")
    
    selected_index = gr.State(None)
    
    with gr.Row():
        with gr.Column():
            generation_credits_display = gr.Number(label="Generation Credits", precision=0, interactive=False)
        with gr.Column():
            train_credits_display = gr.Number(label="Training Credits", precision=0, interactive=False)
        with gr.Column():
            gr.Button("Buy Credits 💳", link="/buy_credits")
    

    with gr.Tabs():
        with gr.TabItem('Create'):
            with gr.Row():
                with gr.Column(scale=3):
                    prompt = gr.Textbox(label="Prompt", 
                                        lines=1, 
                                        placeholder="Enter Your Prompt to start creating 📷", 
                                        info='Some public models may experience longer processing times due to server availability and queue management.')
                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(scale=4):
                    result = gr.Image(label="Imagen Generada")
                
                with gr.Column(scale=3):
                    with gr.Accordion("Public Models"):
                        selected_info = gr.Markdown("")
                        gallery = gr.Gallery(
                            [(item["image_url"], item["model_name"]) for item in LORA_MODELS],
                            label="Public Models",
                            allow_preview=False,
                            columns=3,
                            elem_id="gallery"
                        )
                        
                    with gr.Accordion("Your Models"):
                        user_model_gallery = gr.Gallery(
                            label="Galeria de Modelos",
                            allow_preview=False,
                            columns=3,
                            elem_id="galley"
                        )
                        
            gallery_type = gr.State("Public")

            with gr.Accordion("Advanced Settings", open=False):
                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")
                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=0.95)

            gallery.select(
                update_selection,
                inputs=[gr.State("public"), width, height],
                outputs=[prompt, selected_info, selected_index, width, height, gallery_type]
            )
            
            user_model_gallery.select(
                update_selection,
                inputs=[gr.State("user"), width, height],
                outputs=[prompt, selected_info, selected_index, width, height, gallery_type]
            )
            gr.on(
                triggers=[generate_button.click, prompt.submit],
                fn=run_lora,
                inputs=[prompt, cfg_scale, steps, selected_index, gallery_type, width, height, lora_scale],
                outputs=[result, generation_credits_display]
            )

        with gr.TabItem("Train"):
            gr.Markdown("# Train your own model 🧠")
            gr.Markdown("In this section, you can train your own model using your images.")
            with gr.Row():
                with gr.Column():
                    train_dataset = gr.Gallery(columns=4, interactive=True, label="Tus Imagenes")
                    model_name = gr.Textbox(label="Model Name",)
                    trigger_word = gr.Textbox(label="Trigger Word", 
                                              info="This will be a keyword to later instruct the model when to use these new capabilities we're going to teach it", 
                                              )
                    train_button = gr.Button("Start Training")
            with gr.Accordion("Advanced Settings", open=False):
                train_steps = gr.Slider(label="Training Steps", minimum=100, maximum=10000, step=100, value=1000)
                lora_rank = gr.Number(label='lora_rank', value=16)
                batch_size = gr.Number(label='batch_size', value=1)
                learning_rate = gr.Number(label='learning_rate', value=0.0004)
                training_status = gr.Textbox(label="Training Status")
                
            train_button.click(
                compress_and_train,
                inputs=[train_dataset, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate],
                outputs=[training_status,train_credits_display]
            )
                
        main_demo.load(load_user_models, None, user_model_gallery)
        main_demo.load(load_greet_and_credits, None, [greetings, generation_credits_display, train_credits_display])