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 = r"B:\FLUX\train\output"
TRAIN_DIR = r"B:\FLUX\train"
IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png")
MAX_IMAGES = 150
DESTINATION_PATH = r"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"

def log(message):
    print(f"[Hex Argus Log] {message}")
# Define the TrainPrepClass
class TrainPrepClass:
    def __init__(self):
        pass

    def create_dataset(self, *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(self, images, concept_sentence, *captions):
        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):
                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 get_folder_names(self):
        # 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(self, 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(self, 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(self, 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 = self.load_config(folder_path)
        yaml_content = yaml.dump(config) if config else "Config file not found."
        image_files = self.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_folders(self, folder_names):
        # Delete the selected folders
        for folder_name in folder_names:
            folder_path = os.path.join(OUTPUT_DIR, folder_name)
            if os.path.exists(folder_path):
                shutil.rmtree(folder_path)
        folder_names = self.get_folder_names()
        return gr.update(choices=folder_names, value=[])

    def copy_latest_safetensors_from_folders(self, selected_folders):
        # Copy the latest .safetensors file from each selected folder to the destination folder
        messages = []
        for folder_name in selected_folders:
            folder_path = os.path.join(OUTPUT_DIR, folder_name)
            safetensors_files = [f for f in os.listdir(folder_path) if f.endswith('.safetensors')]
            if safetensors_files:
                latest_file = max(safetensors_files, key=lambda f: os.path.getctime(os.path.join(folder_path, f)))
                source_path = os.path.join(folder_path, latest_file)
                destination_path = os.path.join(DESTINATION_PATH, latest_file)
                shutil.copy(source_path, destination_path)
                messages.append(f"Latest .safetensors file ({latest_file}) copied from {folder_name} to {DESTINATION_PATH}")
            else:
                messages.append(f"No .safetensors files found in {folder_name}")

        return "\n".join(messages)

    def load_captioning(self, 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
        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]
                if base_name in txt_files_dict:
                    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

train_prep = TrainPrepClass()

class QueueingClass:
    def __init__(self):
        self.training_queue = queue.Queue()
        self.start_worker()

    def log(self, message):
        print(f"[QueueingClass Log] {message}")

    def start_worker(self):
        worker_thread = threading.Thread(target=self.worker, daemon=True)
        worker_thread.start()

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

    def add_tasks(self, tasks):
        for task in tasks:
            self.training_queue.put(task)
            self.log("Task has been added to the queue.")

    def get_queue_size(self):
        return self.training_queue.qsize()

queueing_class = QueueingClass()

class TrainingClass:
    def __init__(self):
        self.fake_train = True

    def log(self, message):
        print(f"[TrainingClass Log] {message}")

    def start_training(
        self,
        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 self.fake_train:
            self.log("Fake training enabled. Printing out all project files needed for real training...")
            output_config_file = os.path.join("output", lora_name, "config.yaml")
            if not os.path.exists(output_config_file):
                self.log(f"Configuration file {output_config_file} does not exist.")
            else:
                with open(output_config_file, "r") as f:
                    output_config = yaml.safe_load(f)
                self.log(f"Configuration file content: {yaml.dump(output_config)}")

            config_file = "ai-toolkit/config/examples/train_lora_flux_24gb.yaml"
            if not os.path.exists(config_file):
                self.log(f"Configuration file {config_file} does not exist.")
            else:
                with open(config_file, "r") as f:
                    config = yaml.safe_load(f)
                self.log(f"Configuration file content: {yaml.dump(config)}")

            self.log("Fake training completed.")
            return "Fake training completed successfully."

        output_config_file = os.path.join("output", lora_name, "config.yaml")
        if not os.path.exists(output_config_file):
            self.log(f"Task failed with error: 'Configuration file {output_config_file} does not exist.'")
            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)

        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)

        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)

        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.")
        self.log(f"Starting job for {lora_name} with config: {config_path}")
        job = get_job(config_path)
        if job:
            self.log(f"Job created successfully for {lora_name}")
            job.run()
            self.log(f"Job run completed for {lora_name}")
            job.cleanup()
            self.log(f"Job cleanup completed for {lora_name}")
        else:
            self.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 load_config(self, folder_path):
        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 start_training_task(
        self,
        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]

        if isinstance(lora_name, list):
            lora_name = lora_name[0]

        def task():
            for folder in folder_names:
                print(f"Processing folder: {folder}")
                if folder:
                    folder_path = os.path.join(OUTPUT_DIR, folder)
                    config = self.load_config(folder_path)
                    if not config:
                        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
                                        }
                                    }
                                ]
                            }
                        }
                        config["config"]["process"][0]["trigger_word"] = folder[:3]

                    random_config_name = str(uuid.uuid4())
                    os.makedirs("tmp", exist_ok=True)
                    config_path = f"tmp/{random_config_name}-{folder}.yaml"

                    with open(config_path, "w") as f:
                        yaml.dump(config, f)

                    # Create project in the output folder and copy the file there
                    project_output_path = os.path.join(OUTPUT_DIR, folder)
                    os.makedirs(project_output_path, exist_ok=True)
                    shutil.copy(config_path, os.path.join(project_output_path, "config.yaml"))

                    job = get_job(config_path)
                    job.run()
                    job.cleanup()

                    self.log(f"Training for {folder} completed successfully.")

            return self.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:
            self.log(f"Queuing task for folders: {', '.join(folder_names)}")
        else:
            self.log(f"Queuing task for LoRA: {lora_name}")

        random_config_name = str(uuid.uuid4())
        slugged_lora_name = slugify(lora_name)

        config = {
            "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
                        }
                    }
                ]
            }
        }

        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)

        # Create project in the output folder and copy the file there
        project_output_path = os.path.join(OUTPUT_DIR, slugged_lora_name)
        os.makedirs(project_output_path, exist_ok=True)
        shutil.copy(config_path, os.path.join(project_output_path, "config.yaml"))

        queueing_class.add_task(task)
        return f"Task for {lora_name} has been queued."

training_class = TrainingClass()

# OLD CODE
def hide_captioning():
    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
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 = []
                        captioning_rows = []
                        images_list = []
                        captions_list = []
                        for j in range(1, MAX_IMAGES + 1):
                            captioning_row = gr.Row(visible=False)
                            with captioning_row:
                                image = 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,
                                )
                                caption = gr.Textbox(
                                    label=f"Caption {j}", scale=15, interactive=True
                                )

                            output_components.append(captioning_row)
                            output_components.append(image)
                            output_components.append(caption)
                            caption_list.append(caption)
                            captioning_rows.append(captioning_row)
                            images_list.append(image)
                            captions_list.append(caption)

        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(
        train_prep.load_captioning,
        inputs=[images, concept_sentence],
        outputs=output_components
    )

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

    images.clear(
        hide_captioning,
        outputs=[captioning_area, sample, start]
    )

    start.click(fn=train_prep.create_dataset, inputs=[images] + caption_list, outputs=dataset_folder).then(
        fn=training_class.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=train_prep.run_captioning, inputs=[images, concept_sentence] + caption_list, outputs=caption_list)

    return main_ui

# OLD CODE
order = gr.Textbox(
    label="Order of Execution",
    placeholder="Enter numbers separated by commas, e.g., 1,2,3",
    value="1,2,3",
    interactive=True,
)

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

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

# OLD CODE
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"""
        )
        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
def scan_tab():
    with gr.Tab("Scan"):
        folder_names = train_prep.get_folder_names()

        print("Projects:")
        for project in folder_names:
            print(project)

        # Create CheckboxGroup for projects
        selected_folders = gr.CheckboxGroup(
            choices=[(folder_name, folder_name) for folder_name in folder_names],
            label="Projects"
        )

        delete_button = gr.Button("Delete Selected Folders")
        run_button = gr.Button("Run Selected Projects")
        copy_button = gr.Button("Copy Latest .safetensors from Selected Folders")
        refresh_button = gr.Button("Refresh Projects")
        folder_name_output = gr.Markdown()

        def update_ui(folder_names):
            concatenated_folder_names = ", ".join(folder_names)
            if not concatenated_folder_names:
                return gr.Markdown("No folder selected")
            
            folder_path = os.path.join(OUTPUT_DIR, folder_names[0])
            config = train_prep.load_config(folder_path)
            yaml_content = yaml.dump(config) if config else "Config file not found."
            return gr.Markdown(f"# {concatenated_folder_names}")

        def delete_folders(folder_names):
            for folder_name in folder_names:
                folder_path = os.path.join(OUTPUT_DIR, folder_name)
                if os.path.exists(folder_path):
                    shutil.rmtree(folder_path)
            updated_folder_names = train_prep.get_folder_names()
            return gr.update(choices=updated_folder_names, value=[])

        def copy_latest_safetensors_from_folders(folder_names):
            for folder_name in folder_names:
                folder_path = os.path.join(OUTPUT_DIR, folder_name)
                train_prep.copy_latest_safetensors(folder_path)
            return f"Latest .safetensors files copied to {DESTINATION_PATH}"

        def refresh_projects():
            folder_names = train_prep.get_folder_names()
            return gr.update(choices=folder_names, value=None)

        def start_training_tasks(folder_names):
            for folder_name in folder_names:
                def task():
                    training_class.start_training_task(folder_name=folder_name)
                queueing_class.add_task(task)
            return f"Training tasks for {', '.join(folder_names)} have been queued."

        selected_folders.change(update_ui, inputs=selected_folders, outputs=[folder_name_output])
        delete_button.click(delete_folders, inputs=selected_folders, outputs=selected_folders)
        run_button.click(start_training_tasks, inputs=selected_folders, outputs=folder_name_output)
        copy_button.click(copy_latest_safetensors_from_folders, inputs=selected_folders, outputs=folder_name_output)
        refresh_button.click(lambda: gr.update(choices=train_prep.get_folder_names(), value=[]), inputs=None, outputs=selected_folders)

    return scan_tab

# OLD CODE
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
'''

# OLD CODE
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}
"""

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

# OLD CODE
demo.launch(server_name="0.0.0.0", server_port=7866, show_error=True)

--------

wat is it about our tabs and the classes not communiating correctly