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 queue
import random
import re

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"
FINAL_MODEL_PATH = "models/{slugged_lora_name}.safetensors"
CONFIG_FILE = "ai-toolkit/config/examples/train_lora_flux_24gb.yaml"
JSONL_FILE_PATH = os.path.join("{destination_folder}", "metadata.jsonl")
FULL_DATASET_PATH = os.path.join(TRAIN_DIR, "{dataset_path}")
CONFIG_PATH = "tmp/{random_config_name}-{slugged_lora_name}.yaml"

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

# Create a queue to hold training tasks
training_queue = queue.Queue()

def log(message):
    print(f"[Hex Argus Log] {message}")

def worker():
    while True:
        task = training_queue.get()  
        try:
            task()  
        except Exception as e:
            log(f"Task failed with error: {str(e)}")
        finally:
            training_queue.task_done()

worker_thread = threading.Thread(target=worker, daemon=True)
worker_thread.start()

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 = FULL_DATASET_PATH.format(dataset_path=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):
    if isinstance(folder_name, list) and folder_name:  # Ensure folder_name is a non-empty list
        folder_name = folder_name[0]  # Ensure folder_name is a string
    elif not folder_name:
        return gr.Markdown("No folder selected"), gr.Code("No config available", language="yaml"), gr.Gallery([])

    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(folder_name):
    # Delete the selected folder
    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)
import os
import uuid
import yaml
import logging
import gradio as gr
from queue import Queue

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
training_queue = Queue()

def start_training_task(
    lora_name=None,
    concept_sentence=None,
    steps=None,
    lr=None,
    rank=None,
    model_to_train=None,
    low_vram=None,
    dataset_folder=None,
    sample_1=None,
    sample_2=None,
    sample_3=None,
    use_more_advanced_options=None,
    more_advanced_options=None,
    push_to_hub=None,
    folder_name=None
):
    folder_names = folder_name if isinstance(folder_name, list) else [folder_name] if folder_name else []
    lora_name = lora_name[0] if isinstance(lora_name, list) else lora_name

    def task():
        try:
            for folder in folder_names:
                if not folder:
                    continue
                logger.info(f"Processing folder: {folder}")
                folder_path = os.path.join(OUTPUT_DIR, folder)
                config = load_config(folder_path)
                if not config:
                    raise gr.Error("Config file not found.")
                config_name = f"{uuid.uuid4()}-{folder}.yaml"
                os.makedirs(folder_path, exist_ok=True)
                config_path = os.path.join(folder_path, config_name)
                with open(config_path, "w") as f:
                    yaml.dump(config, f)
                job = get_job(config_path)
                job.run()
                job.cleanup()
                os.remove(config_path)
                logger.info(f"Training for {folder} completed successfully.")
            return 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,
            )
        except Exception as e:
            logger.error(f"An error occurred: {e}")
            raise

    if folder_names:
        logger.info(f"Queuing task for folders: {', '.join(folder_names)}")
    else:
        logger.info(f"Queuing task for LoRA: {lora_name}")
    training_queue.put(task)
    return "Task has been queued."

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

def extract_folder_path(yaml_content):
    match = re.search(r'folder_path:\s*(\S+)', yaml_content)
    if match:
        return match.group(1)
    return None
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,
):
    # Load the relevant YAML file from the output folder
    output_config_file = os.path.join("output", lora_name, "config.yaml")
    if not os.path.exists(output_config_file):
        log(f"[Hex Argus Log] Task failed with error: 'Configuration file {output_config_file} does not exist.'")
        # Create a default configuration if not found
        output_config = {
            "config": {
                "name": lora_name,
                "process": [
                    {
                        "model": {
                            "name_or_path": "default-model",
                            "low_vram": low_vram
                        },
                        "train": {
                            "steps": steps,
                            "lr": lr,
                            "skip_first_sample": True
                        },
                        "network": {
                            "linear": rank,
                            "linear_alpha": rank
                        },
                        "datasets": [
                            {
                                "folder_path": dataset_folder
                            }
                        ],
                        "save": {
                            "push_to_hub": push_to_hub
                        },
                        "sample": {
                            "sample_every": steps,
                            "sample_steps": 28,
                            "prompts": [sample_1, sample_2, sample_3]
                        }
                    }
                ]
            }
        }
    else:
        with open(output_config_file, "r") as f:
            output_config = yaml.safe_load(f)
    
    # Use the values from the output YAML file for all relevant keys
    config_file = "ai-toolkit/config/examples/train_lora_flux_24gb.yaml"
    if not os.path.exists(config_file):
        raise gr.Error(f"Configuration file {config_file} does not exist.")
    
    with open(config_file, "r") as f:
        config = yaml.safe_load(f)
    
    # Update the training configuration with values from the output YAML file
    config["config"]["name"] = output_config["config"]["name"]
    config["config"]["process"][0]["model"] = output_config["config"]["process"][0]["model"]
    config["config"]["process"][0]["train"] = output_config["config"]["process"][0]["train"]
    config["config"]["process"][0]["network"] = output_config["config"]["process"][0]["network"]
    config["config"]["process"][0]["datasets"] = output_config["config"]["process"][0]["datasets"]
    config["config"]["process"][0]["save"] = output_config["config"]["process"][0]["save"]
    config["config"]["process"][0]["sample"] = output_config["config"]["process"][0]["sample"]
    
    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)
    
    # Copy the config.yaml from the output/{project name} folder
    config_path = os.path.join("output", lora_name, "config.yaml")
    if not os.path.exists(config_path):
        raise gr.Error(f"Configuration file {config_path} does not exist.")
    # Start the training job with the updated configuration
    log(f"Starting job for {lora_name} with config: {config_path}")
    job = get_job(config_path)
    if job:
        log(f"Job created successfully for {lora_name}")
        job.run()
        log(f"Job run completed for {lora_name}")
        job.cleanup()
        log(f"Job cleanup completed for {lora_name}")
    else:
        log(f"Failed to create job for {lora_name}")

    final_model_path = f"models/{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 {lora_name} and copied to {destination_path}"

def start_training_task(
    lora_name=None,
    concept_sentence=None,
    steps=None,
    lr=None,
    rank=None,
    model_to_train=None,
    low_vram=None,
    dataset_folder=None,
    sample_1=None,
    sample_2=None,
    sample_3=None,
    use_more_advanced_options=None,
    more_advanced_options=None,
    push_to_hub=None,
    folder_name=None
):
    if isinstance(folder_name, list):
        folder_names = folder_name
    else:
        folder_names = [folder_name]
    
    # Ensure lora_name is a string
    if isinstance(lora_name, list):
        lora_name = lora_name[0]

    def task():
        for folder in folder_names:
            print(f"Processing folder: {folder}")
            if folder:
                # Run the selected project based on its configuration
                folder_path = os.path.join(OUTPUT_DIR, folder)
                config = load_config(folder_path)
                if not config:
                    # Create a default configuration if not found
                    config = {
                        "config": {
                            "name": folder,
                            "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}"
                                        }
                                    ],
                                    "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[:3]
                
                # Generate a unique name for the config
                random_config_name = str(uuid.uuid4())
                os.makedirs("tmp", exist_ok=True)
                config_path = f"tmp/{random_config_name}-{folder}.yaml"
                
                # Save the config to a temporary file
                with open(config_path, "w") as f:
                    yaml.dump(config, f)
                
                # Queue the task (if config is available)
                job = get_job(config_path)
                job.run()
                job.cleanup()

                log(f"Training for {folder} completed successfully.")
        
        # Start training with provided parameters
        return 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 folder_name:
        log(f"Queuing task for folders: {', '.join(folder_names)}")
    else:
        log(f"Queuing task for LoRA: {lora_name}")
    
    # Create and save the configuration file before adding to the queue
    random_config_name = str(uuid.uuid4())
    slugged_lora_name = slugify(lora_name)

    config = {
        # Sample config structure based on user input
        "config": {
            "name": slugged_lora_name,
            "process": [
                {
                    "model": {
                        "name_or_path": model_to_train,
                        "low_vram": low_vram
                    },
                    "train": {
                        "steps": steps,
                        "lr": lr
                    },
                    "network": {
                        "linear": rank,
                        "linear_alpha": rank
                    },
                    "datasets": [
                        {"folder_path": dataset_folder}
                    ],
                    "save": {
                        "push_to_hub": push_to_hub
                    }
                }
            ]
        }
    }

    # Save the config as a YAML file
    config_path = f"tmp/{random_config_name}-{slugged_lora_name}.yaml"
    os.makedirs("tmp", exist_ok=True)
    with open(config_path, "w") as f:
        yaml.dump(config, f)

    # Add the task to the queue
    training_queue.put(task)
    return f"Task for {lora_name} has been queued."

def scan_tab():
    with gr.Tab("Scan"):
        folder_names = get_folder_names()
        
        print("Projects:")
        for project in folder_names:
            print(project)
        
        # Create CheckboxGroup for projects
        selected_folder = gr.CheckboxGroup(
            choices=[(folder_name, folder_name) for folder_name in folder_names], 
            label="Projects"
        )
        
        delete_button = gr.Button("Delete Folder")
        run_button = gr.Button("Run Selected Project")
        copy_button = gr.Button("Copy Latest .safetensors")
        folder_name_output = gr.Markdown()
        
        yaml_output = gr.Code()
        image_gallery = gr.Gallery()
        
        selected_folder.change(update_ui, inputs=selected_folder, outputs=[folder_name_output, yaml_output, image_gallery])
        
        delete_button.click(delete_folder, inputs=selected_folder, outputs=selected_folder)
        run_button.click(start_training_task, inputs=selected_folder, outputs=folder_name_output)
        copy_button.click(copy_latest_safetensors, inputs=selected_folder, outputs=folder_name_output)
    
    return scan_tab

def copy_latest_safetensors(selected_folder):
    # Copy the latest .safetensors file to the destination folder
    folder_name = selected_folder
    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.")
    
    latest_file = max(safetensors_files, key=os.path.getctime)
    source_path = os.path.join(folder_path, latest_file)
    destination_path = os.path.join(DESTINATION_PATH, latest_file)
    shutil.copy(source_path, destination_path)
    
    return f"Latest .safetensors file ({latest_file}) copied to {DESTINATION_PATH}"

def create_ui_component(i):
    with gr.Column() as main_ui:
        with gr.Row():
            lora_name = gr.Textbox(
                label=f"LoRA Name {i+1}",
                info="This has to be a unique name",
                placeholder="e.g.: Persian Miniature Painting style, Cat Toy",
            )
            concept_sentence = gr.Textbox(
                label=f"Trigger {i+1}",
                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=f"Images {i+1}",
                    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(f"Add AI captions {i+1}")
                        output_components = [captioning_area]
                        caption_list = []
                        for j in range(1, MAX_IMAGES + 1):
                            locals()[f"captioning_row_{i}_{j}"] = gr.Row(visible=False)
                            with locals()[f"captioning_row_{i}_{j}"]:
                                locals()[f"image_{i}_{j}"] = 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}_{j}"] = gr.Textbox(
                                    label=f"Caption {j}", scale=15, interactive=True
                                )

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

        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(f"Start training {i+1}", 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]
    )
    
    def start_training_wrapper(*args):
        while True:
            if not any(t.is_alive() for t in threading.enumerate() if t.name.startswith("training_thread_")):
                break
            time.sleep(1)
        return start_training(*args)

    start.click(fn=create_dataset, inputs=[images] + caption_list, outputs=dataset_folder).then(
        fn=start_training_task,
        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
order = gr.Textbox(
    label="Order of Execution",
    placeholder="Enter numbers separated by commas, e.g., 1,2,3",
    value="1,2,3",
    interactive=True,
)

queue_status = gr.Markdown(
    label="Queue Status",
    value="Queue is empty"
)

def update_queue_status():
    while True:
        queue_size = training_queue.qsize()
        if queue_size == 0:
            queue_status.value = "Queue is empty"
        else:
            queue_status.value = f"Queue has {queue_size} task(s)"
        time.sleep(1)

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.Row():
            ui_components = [create_ui_component(i) for i in range(3)]
            order.change(
                fn=lambda x: sorted(ui_components, key=lambda i: int(x.split(",")[i])),
                inputs=order,
                outputs=ui_components
            )
        with gr.Row():
            queue_status.render()
            
    threading.Thread(target=update_queue_status, daemon=True).start()
    return flux_train_tab

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}
"""

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)


------

give me all the training related code