import os
import sys
import uuid
import shutil
import json
import yaml
from datetime import datetime
from slugify import slugify
from PIL import Image
import time
import threading

import gradio as gr
import torch
from huggingface_hub import whoami
from transformers import AutoProcessor, AutoModelForCausalLM

# Add the current working directory to the Python path
sys.path.insert(0, os.getcwd())
sys.path.insert(0, "ai-toolkit")

from toolkit.job import get_job

# Configuration
OUTPUT_DIR = "B:\\FLUX\\train\\output"
TRAIN_DIR = "B:\\FLUX\\train"
IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png")
MAX_IMAGES = 150
DESTINATION_PATH = "X:\\AI\\AI-Resources\\models\\lora"

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"

def get_folder_names():
    # Get folder names sorted by modification time
    folders = [(folder, os.path.getmtime(os.path.join(OUTPUT_DIR, folder))) for folder in os.listdir(OUTPUT_DIR) if os.path.isdir(os.path.join(OUTPUT_DIR, folder))]
    folders.sort(key=lambda x: x[1], reverse=True)
    return [folder[0] for folder in folders]

def load_config(folder_path):
    # Load configuration from YAML file
    config_path = os.path.join(folder_path, "config.yaml")
    if os.path.exists(config_path):
        with open(config_path, "r") as file:
            return yaml.safe_load(file)
    return None

def get_image_files(config):
    # Retrieve image files based on the configuration
    image_files = []
    for process in config.get("config", {}).get("process", []):
        for dataset in process.get("datasets", []):
            dataset_path = dataset.get("folder_path")
            if dataset_path:
                full_dataset_path = os.path.join(TRAIN_DIR, dataset_path)
                if os.path.exists(full_dataset_path):
                    image_files.extend([os.path.join(full_dataset_path, f) for f in os.listdir(full_dataset_path) if f.endswith(IMAGE_EXTENSIONS)])
    return image_files

def update_ui(folder_name):
    # Update UI components based on the selected folder
    folder_path = os.path.join(OUTPUT_DIR, folder_name)
    config = load_config(folder_path)
    yaml_content = yaml.dump(config) if config else "Config file not found."
    image_files = get_image_files(config) if config else []
    return gr.Markdown(f"# {folder_name}"), gr.Code(yaml_content, language="yaml"), gr.Gallery(image_files)

def delete_folder(selected_folder_with_safetensors, selected_folder_without_safetensors, selected_folder_missing_files):
    # Delete the selected folder
    folder_name = selected_folder_with_safetensors or selected_folder_without_safetensors or selected_folder_missing_files
    if folder_name:
        folder_path = os.path.join(OUTPUT_DIR, folder_name)
        if os.path.exists(folder_path):
            shutil.rmtree(folder_path)
    folder_names = get_folder_names()
    return gr.update(choices=folder_names, value=None), gr.update(choices=folder_names, value=None), gr.update(choices=folder_names, value=None)

def run_selected_project(folder_name):
    folder_path = os.path.join(OUTPUT_DIR, folder_name)
    config = load_config(folder_path)
    
    if not config:
        # Create a default configuration if not found
        config = {
            "config": {
                "name": folder_name,
                "process": [
                    {
                        "model": {
                            "name_or_path": "default-model",
                            "low_vram": False
                        },
                        "train": {
                            "steps": 1000,
                            "lr": 0.0001,
                            "skip_first_sample": True
                        },
                        "network": {
                            "linear": 16,
                            "linear_alpha": 16
                        },
                        "datasets": [
                            {
                                "folder_path": f"datasets/{folder_name}"
                            }
                        ],
                        "save": {
                            "push_to_hub": False
                        }
                    }
                ]
            }
        }
        # Use the first 3 characters of the folder name as the trigger word
        config["config"]["process"][0]["trigger_word"] = folder_name[:3]
    
    random_config_name = str(uuid.uuid4())
    os.makedirs("tmp", exist_ok=True)
    config_path = f"tmp/{random_config_name}-{folder_name}.yaml"
    
    with open(config_path, "w") as f:
        yaml.dump(config, f)
    
    job = get_job(config_path)
    job.run()
    job.cleanup()
    
    return f"Training for {folder_name} completed successfully."

def update_folder_lists():
    # Update the list of folders with and without .safetensors files
    folder_names = get_folder_names()
    projects_with_safetensors = []
    projects_without_safetensors = []
    projects_missing_files = []

    for folder in folder_names:
        folder_path = os.path.join(OUTPUT_DIR, folder)
        safetensors_files = [f for f in os.listdir(folder_path) if f.endswith(".safetensors")]
        if safetensors_files:
            projects_with_safetensors.append(folder)
        else:
            projects_without_safetensors.append(folder)

    # Check for dataset folders that need YAML and other project files
    dataset_folders = [f for f in os.listdir(TRAIN_DIR) if os.path.isdir(os.path.join(TRAIN_DIR, f))]
    for dataset_folder in dataset_folders:
        dataset_folder_path = os.path.join(TRAIN_DIR, dataset_folder)
        if not os.path.exists(os.path.join(dataset_folder_path, "config.yaml")):
            projects_missing_files.append(dataset_folder)
        else:
            # Check if the dataset folder is referenced in any project config
            referenced = False
            for project_folder in folder_names:
                project_config_path = os.path.join(OUTPUT_DIR, project_folder, "config.yaml")
                if os.path.exists(project_config_path):
                    with open(project_config_path, "r") as file:
                        config = yaml.safe_load(file)
                        for process in config.get("config", {}).get("process", []):
                            for dataset in process.get("datasets", []):
                                if dataset.get("folder_path") == f"datasets/{dataset_folder}":
                                    referenced = True
                                    break
                if referenced:
                    break
            if not referenced:
                projects_missing_files.append(dataset_folder)

    return projects_with_safetensors, projects_without_safetensors, projects_missing_files

def periodic_update(selected_folder_with_safetensors, selected_folder_without_safetensors, selected_folder_missing_files):
    # Periodically update the folder lists
    while True:
        projects_with_safetensors, projects_without_safetensors, projects_missing_files = update_folder_lists()
        selected_folder_with_safetensors.update(choices=projects_with_safetensors, value=selected_folder_with_safetensors.value)
        selected_folder_without_safetensors.update(choices=projects_without_safetensors, value=selected_folder_without_safetensors.value)
        selected_folder_missing_files.update(choices=projects_missing_files, value=selected_folder_missing_files.value)
        time.sleep(5)

def copy_latest_safetensors(selected_folder_with_safetensors, selected_folder_without_safetensors):
    # Copy the latest .safetensors file to the destination folder
    folder_name = selected_folder_with_safetensors or selected_folder_without_safetensors
    if not folder_name:
        raise gr.Error("No folder selected.")
    
    folder_path = os.path.join(OUTPUT_DIR, folder_name)
    safetensors_files = [f for f in os.listdir(folder_path) if f.endswith(".safetensors")]
    if not safetensors_files:
        raise gr.Error("No .safetensors files found in the selected folder.")
    
    # Get the latest .safetensors file
    latest_file = max(safetensors_files, key=lambda f: os.path.getmtime(os.path.join(folder_path, f)))
    latest_file_path = os.path.join(folder_path, latest_file)
    
    # Copy the file to the destination folder
    shutil.copy(latest_file_path, DESTINATION_PATH)
    
    return f"Copied {latest_file} to {DESTINATION_PATH}"

def scan_tab():
    with gr.Tab("Scan"):
        projects_with_safetensors, projects_without_safetensors, projects_missing_files = update_folder_lists()
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Projects with .safetensors files")
                selected_folder_with_safetensors = gr.Radio(projects_with_safetensors, label="Select Folder")
            with gr.Column():
                gr.Markdown("### Projects without .safetensors files")
                selected_folder_without_safetensors = gr.Radio(projects_without_safetensors, label="Select Folder")
            with gr.Column():
                gr.Markdown("### Dataset folders missing files")
                selected_folder_missing_files = gr.Radio(projects_missing_files, label="Select Folder")
        
        delete_button = gr.Button("Delete Folder")
        run_button = gr.Button("Run Selected Project")
        copy_button = gr.Button("Copy Latest .safetensors to Destination")
        folder_name_output = gr.Markdown()
        
        with gr.Row():
            with gr.Column():
                with gr.Accordion("Config", open=False):
                    yaml_output = gr.Code()
            with gr.Column():
                with gr.Accordion("Images", open=False):
                    image_gallery = gr.Gallery()
        
        selected_folder_with_safetensors.change(
            update_ui, 
            inputs=selected_folder_with_safetensors, 
            outputs=[folder_name_output, yaml_output, image_gallery]
        ).then(
            lambda: gr.update(value=None), 
            inputs=None, 
            outputs=selected_folder_without_safetensors
        ).then(
            lambda: gr.update(value=None), 
            inputs=None, 
            outputs=selected_folder_missing_files
        )
        
        selected_folder_without_safetensors.change(
            update_ui, 
            inputs=selected_folder_without_safetensors, 
            outputs=[folder_name_output, yaml_output, image_gallery]
        ).then(
            lambda: gr.update(value=None), 
            inputs=None, 
            outputs=selected_folder_with_safetensors
        ).then(
            lambda: gr.update(value=None), 
            inputs=None, 
            outputs=selected_folder_missing_files
        )
        
        selected_folder_missing_files.change(
            update_ui, 
            inputs=selected_folder_missing_files, 
            outputs=[folder_name_output, yaml_output, image_gallery]
        ).then(
            lambda: gr.update(value=None), 
            inputs=None, 
            outputs=selected_folder_with_safetensors
        ).then(
            lambda: gr.update(value=None), 
            inputs=None, 
            outputs=selected_folder_without_safetensors
        )
        
        delete_button.click(
            delete_folder, 
            inputs=[selected_folder_with_safetensors, selected_folder_without_safetensors, selected_folder_missing_files], 
            outputs=[selected_folder_with_safetensors, selected_folder_without_safetensors, selected_folder_missing_files]
        )
        
        run_button.click(
            run_selected_project, 
            inputs=selected_folder_without_safetensors, 
            outputs=folder_name_output
        ).then(
            run_selected_project, 
            inputs=selected_folder_missing_files, 
            outputs=folder_name_output
        )
        
        copy_button.click(
            copy_latest_safetensors, 
            inputs=[selected_folder_with_safetensors, selected_folder_without_safetensors], 
            outputs=folder_name_output
        )
    
    # Start the periodic update in a separate thread
    threading.Thread(target=periodic_update, args=(selected_folder_with_safetensors, selected_folder_without_safetensors, selected_folder_missing_files), daemon=True).start()
    
    return scan_tab

def load_captioning(uploaded_files, concept_sentence):
    uploaded_images = [file for file in uploaded_files if not file.endswith('.txt')]
    txt_files = [file for file in uploaded_files if file.endswith('.txt')]
    txt_files_dict = {os.path.splitext(os.path.basename(txt_file))[0]: txt_file for txt_file in txt_files}
    updates = []
    if len(uploaded_images) <= 1:
        raise gr.Error(
            "Please upload at least 2 images to train your model (the ideal number with default settings is between 4-30)"
        )
    elif len(uploaded_images) > MAX_IMAGES:
        raise gr.Error(f"For now, only {MAX_IMAGES} or less images are allowed for training")
    # Update for the captioning_area
    # for _ in range(3):
    updates.append(gr.update(visible=True))
    # Update visibility and image for each captioning row and image
    for i in range(1, MAX_IMAGES + 1):
        # Determine if the current row and image should be visible
        visible = i <= len(uploaded_images)
        
        # Update visibility of the captioning row
        updates.append(gr.update(visible=visible))

        # Update for image component - display image if available, otherwise hide
        image_value = uploaded_images[i - 1] if visible else None
        updates.append(gr.update(value=image_value, visible=visible))
        
        corresponding_caption = False
        if(image_value):
            base_name = os.path.splitext(os.path.basename(image_value))[0]
            print(base_name)
            print(image_value)
            if base_name in txt_files_dict:
                print("entrou")
                with open(txt_files_dict[base_name], 'r') as file:
                    corresponding_caption = file.read()
                    
        # Update value of captioning area
        text_value = corresponding_caption if visible and corresponding_caption else "[trigger]" if visible and concept_sentence else None
        updates.append(gr.update(value=text_value, visible=visible))

    # Update for the sample caption area
    updates.append(gr.update(visible=True))
    # Update prompt samples
    updates.append(gr.update(placeholder=f'A portrait of person in a bustling cafe {concept_sentence}', value=f'A person in a bustling cafe {concept_sentence}'))
    updates.append(gr.update(placeholder=f"A mountainous landscape in the style of {concept_sentence}"))
    updates.append(gr.update(placeholder=f"A {concept_sentence} in a mall"))
    updates.append(gr.update(visible=True))
    return updates

def hide_captioning():
    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) 

def create_dataset(*inputs):
    print("Creating dataset")
    images = inputs[0]
    destination_folder = str(f"datasets/{uuid.uuid4()}")
    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)

    jsonl_file_path = os.path.join(destination_folder, "metadata.jsonl")
    with open(jsonl_file_path, "a") as jsonl_file:
        for index, image in enumerate(images):
            new_image_path = shutil.copy(image, destination_folder)

            original_caption = inputs[index + 1]
            file_name = os.path.basename(new_image_path)

            data = {"file_name": file_name, "prompt": original_caption}

            jsonl_file.write(json.dumps(data) + "\n")

    return destination_folder


def run_captioning(images, concept_sentence, *captions):
    #Load internally to not consume resources for training
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch_dtype = torch.float16
    model = AutoModelForCausalLM.from_pretrained(
        "multimodalart/Florence-2-large-no-flash-attn", torch_dtype=torch_dtype, trust_remote_code=True
    ).to(device)
    processor = AutoProcessor.from_pretrained("multimodalart/Florence-2-large-no-flash-attn", trust_remote_code=True)

    captions = list(captions)
    for i, image_path in enumerate(images):
        print(captions[i])
        if isinstance(image_path, str):  # If image is a file path
            image = Image.open(image_path).convert("RGB")

        prompt = "<DETAILED_CAPTION>"
        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)

        generated_ids = model.generate(
            input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, num_beams=3
        )

        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
        parsed_answer = processor.post_process_generation(
            generated_text, task=prompt, image_size=(image.width, image.height)
        )
        caption_text = parsed_answer["<DETAILED_CAPTION>"].replace("The image shows ", "")
        if concept_sentence:
            caption_text = f"{caption_text} [trigger]"
        captions[i] = caption_text

        yield captions
    model.to("cpu")
    del model
    del processor

def recursive_update(d, u):
    for k, v in u.items():
        if isinstance(v, dict) and v:
            d[k] = recursive_update(d.get(k, {}), v)
        else:
            d[k] = v
    return d

import shutil

def start_training(
    lora_name,
    concept_sentence,
    steps,
    lr,
    rank,
    model_to_train,
    low_vram,
    dataset_folder,
    sample_1,
    sample_2,
    sample_3,
    use_more_advanced_options,
    more_advanced_options,
    push_to_hub,
):
    if not lora_name:
        raise gr.Error("You forgot to insert your LoRA name! This name has to be unique.")
    try:
        if push_to_hub and (whoami()["auth"]["accessToken"]["role"] == "write" or "repo.write" in whoami()["auth"]["accessToken"]["fineGrained"]["scoped"][0]["permissions"]):
            gr.Info(f"Starting training locally {whoami()['name']}. Your LoRA will be available locally and in Hugging Face after it finishes.")
        else:
            push_to_hub = False
            gr.Warning("Started training locally. Your LoRa will only be available locally because you didn't login with a `write` token to Hugging Face")
    except:
        push_to_hub = False
        gr.Warning("Started training locally. Your LoRa will only be available locally because you didn't login with a `write` token to Hugging Face")
            
    print("Started training")
    slugged_lora_name = slugify(lora_name)

    # Load the default config
    with open("ai-toolkit/config/examples/train_lora_flux_24gb.yaml", "r") as f:
        config = yaml.safe_load(f)

    # Update the config with user inputs
    config["config"]["name"] = slugged_lora_name
    config["config"]["process"][0]["model"]["low_vram"] = low_vram
    config["config"]["process"][0]["train"]["skip_first_sample"] = True
    config["config"]["process"][0]["train"]["steps"] = int(steps)
    config["config"]["process"][0]["train"]["lr"] = float(lr)
    config["config"]["process"][0]["network"]["linear"] = int(rank)
    config["config"]["process"][0]["network"]["linear_alpha"] = int(rank)
    config["config"]["process"][0]["datasets"][0]["folder_path"] = dataset_folder
    config["config"]["process"][0]["save"]["push_to_hub"] = push_to_hub
    if push_to_hub:
        try:
            username = whoami()["name"]
        except:
            raise gr.Error("Error trying to retrieve your username. Are you sure you are logged in with Hugging Face?")
        config["config"]["process"][0]["save"]["hf_repo_id"] = f"{username}/{slugged_lora_name}"
        config["config"]["process"][0]["save"]["hf_private"] = True
    if concept_sentence:
        config["config"]["process"][0]["trigger_word"] = concept_sentence
    
    if sample_1 or sample_2 or sample_3:
        config["config"]["process"][0]["train"]["disable_sampling"] = False
        config["config"]["process"][0]["sample"]["sample_every"] = steps
        config["config"]["process"][0]["sample"]["sample_steps"] = 28
        config["config"]["process"][0]["sample"]["prompts"] = []
        if sample_1:
            config["config"]["process"][0]["sample"]["prompts"].append(sample_1)
        if sample_2:
            config["config"]["process"][0]["sample"]["prompts"].append(sample_2)
        if sample_3:
            config["config"]["process"][0]["sample"]["prompts"].append(sample_3)
    else:
        config["config"]["process"][0]["train"]["disable_sampling"] = True
    if(model_to_train == "schnell"):
        config["config"]["process"][0]["model"]["name_or_path"] = "black-forest-labs/FLUX.1-schnell"
        config["config"]["process"][0]["model"]["assistant_lora_path"] = "ostris/FLUX.1-schnell-training-adapter"
        config["config"]["process"][0]["sample"]["sample_steps"] = 4
    if(use_more_advanced_options):
        more_advanced_options_dict = yaml.safe_load(more_advanced_options)
        config["config"]["process"][0] = recursive_update(config["config"]["process"][0], more_advanced_options_dict)
        print(config)
    
    # Save the updated config
    # generate a random name for the config
    random_config_name = str(uuid.uuid4())
    os.makedirs("tmp", exist_ok=True)
    config_path = f"tmp/{random_config_name}-{slugged_lora_name}.yaml"
    with open(config_path, "w") as f:
        yaml.dump(config, f)
    
    # run the job locally
    job = get_job(config_path)
    job.run()
    job.cleanup()

    # Copy the final safetensors file to the specified directory
    final_model_path = f"models/{slugged_lora_name}.safetensors"
    destination_path = "X:/AI/AI-Resources/models/lora"
    if os.path.exists(final_model_path):
        shutil.copy(final_model_path, destination_path)

    return f"Training completed successfully. Model saved as {slugged_lora_name} and copied to {destination_path}"

config_yaml = '''
device: cuda:0
model:
  is_flux: true
  quantize: true
network:
  linear: 16 #it will overcome the 'rank' parameter
  linear_alpha: 16 #you can have an alpha different than the ranking if you'd like
  type: lora
sample:
  guidance_scale: 3.5
  height: 1024
  neg: '' #doesn't work for FLUX
  sample_every: 1000
  sample_steps: 28
  sampler: flowmatch
  seed: 42
  walk_seed: true
  width: 1024
save:
  dtype: float16
  hf_private: true
  max_step_saves_to_keep: 4
  push_to_hub: false  # Set to false to disable uploading by default
  save_every: 10000
train:
  batch_size: 1
  dtype: bf16
  ema_config:
    ema_decay: 0.99
    use_ema: true
  gradient_accumulation_steps: 1
  gradient_checkpointing: true
  noise_scheduler: flowmatch 
  optimizer: adamw8bit #options: prodigy, dadaptation, adamw, adamw8bit, lion, lion8bit
  train_text_encoder: false #probably doesn't work for flux
  train_unet: true
'''

theme = gr.themes.Monochrome(
    text_size=gr.themes.Size(lg="18px", md="15px", sm="13px", xl="22px", xs="12px", xxl="24px", xxs="9px"),
    font=[gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui", "sans-serif"],
)
css = """
h1{font-size: 2em}
h3{margin-top: 0}
#component-1{text-align:center}
.main_ui_logged_out{opacity: 0.3; pointer-events: none}
.tabitem{border: 0px}
.group_padding{padding: .55em}
"""
def flux_train_tab():
    with gr.Tab("Flux Train"):
        gr.Markdown(
            """# LoRA Ease for FLUX 🧞‍♂️
    ### Train a high quality FLUX LoRA in a breeze ༄ using [Ostris' AI Toolkit](https://github.com/ostris/ai-toolkit)"""
        )
        with gr.Column() as main_ui:
            with gr.Row():
                lora_name = gr.Textbox(
                    label="The name of your LoRA",
                    info="This has to be a unique name",
                    placeholder="e.g.: Persian Miniature Painting style, Cat Toy",
                )
                concept_sentence = gr.Textbox(
                    label="Trigger word/sentence",
                    info="Trigger word or sentence to be used",
                    placeholder="uncommon word like p3rs0n or trtcrd, or sentence like 'in the style of CNSTLL'",
                    interactive=True,
                )
            with gr.Group(visible=True) as image_upload:
                with gr.Row():
                    images = gr.File(
                        file_types=["image", ".txt"],
                        label="Upload your images",
                        file_count="multiple",
                        interactive=True,
                        visible=True,
                        scale=1,
                    )
                    with gr.Column(scale=3, visible=False) as captioning_area:
                        with gr.Column():
                            gr.Markdown(
                                """# Custom captioning
    <p style="margin-top:0">You can optionally add a custom caption for each image (or use an AI model for this). [trigger] will represent your concept sentence/trigger word.</p>
    """, elem_classes="group_padding")
                            do_captioning = gr.Button("Add AI captions with Florence-2")
                            output_components = [captioning_area]
                            caption_list = []
                            for i in range(1, MAX_IMAGES + 1):
                                locals()[f"captioning_row_{i}"] = gr.Row(visible=False)
                                with locals()[f"captioning_row_{i}"]:
                                    locals()[f"image_{i}"] = gr.Image(
                                        type="filepath",
                                        width=111,
                                        height=111,
                                        min_width=111,
                                        interactive=False,
                                        scale=2,
                                        show_label=False,
                                        show_share_button=False,
                                        show_download_button=False,
                                    )
                                    locals()[f"caption_{i}"] = gr.Textbox(
                                        label=f"Caption {i}", scale=15, interactive=True
                                    )

                                output_components.append(locals()[f"captioning_row_{i}"])
                                output_components.append(locals()[f"image_{i}"])
                                output_components.append(locals()[f"caption_{i}"])
                                caption_list.append(locals()[f"caption_{i}"])

            with gr.Accordion("Advanced options", open=False):
                steps = gr.Number(label="Steps", value=1000, minimum=1, maximum=10000, step=1)
                lr = gr.Number(label="Learning Rate", value=4e-4, minimum=1e-6, maximum=1e-3, step=1e-6)
                rank = gr.Number(label="LoRA Rank", value=16, minimum=4, maximum=128, step=4)
                model_to_train = gr.Radio(["dev", "schnell"], value="dev", label="Model to train")
                low_vram = gr.Checkbox(label="Low VRAM", value=True)
                push_to_hub = gr.Checkbox(label="Push to Hugging Face Hub", value=False)
                with gr.Accordion("Even more advanced options", open=False):
                    use_more_advanced_options = gr.Checkbox(label="Use more advanced options", value=False)
                    more_advanced_options = gr.Code(config_yaml, language="yaml")

            with gr.Accordion("Sample prompts (optional)", visible=False) as sample:
                gr.Markdown(
                    "Include sample prompts to test out your trained model. Don't forget to include your trigger word/sentence (optional)"
                )
                sample_1 = gr.Textbox(label="Test prompt 1")
                sample_2 = gr.Textbox(label="Test prompt 2")
                sample_3 = gr.Textbox(label="Test prompt 3")
            
            output_components.append(sample)
            output_components.append(sample_1)
            output_components.append(sample_2)
            output_components.append(sample_3)
            start = gr.Button("Start training", visible=False)
            output_components.append(start)
            progress_area = gr.Markdown("")

        dataset_folder = gr.State()

        images.upload(
            load_captioning,
            inputs=[images, concept_sentence],
            outputs=output_components
        )
        
        images.delete(
            load_captioning,
            inputs=[images, concept_sentence],
            outputs=output_components
        )

        images.clear(
            hide_captioning,
            outputs=[captioning_area, sample, start]
        )
        
        start.click(fn=create_dataset, inputs=[images] + caption_list, outputs=dataset_folder).then(
            fn=start_training,
            inputs=[
                lora_name,
                concept_sentence,
                steps,
                lr,
                rank,
                model_to_train,
                low_vram,
                dataset_folder,
                sample_1,
                sample_2,
                sample_3,
                use_more_advanced_options,
                more_advanced_options,
                push_to_hub
            ],
            outputs=progress_area,
        )

        do_captioning.click(fn=run_captioning, inputs=[images, concept_sentence] + caption_list, outputs=caption_list)
        
    return main_ui

with gr.Blocks(theme=theme, css=css) as demo:
    flux_train_tab()
    scan_tab()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7866, show_error=True)