import gradio as gr
import torch
import os
import shutil
import requests
import subprocess
from subprocess import getoutput
from huggingface_hub import snapshot_download, HfApi, create_repo

api = HfApi()

hf_token = os.environ.get("HF_TOKEN_WITH_WRITE_PERMISSION")

def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instance_prompt, max_train_steps, checkpoint_steps):
    
    script_filename = "train_dreambooth_b-lora_sdxl.py"  # Assuming it's in the same folder

    command = [
        "accelerate",
        "launch",
        script_filename,  # Use the local script
        "--pretrained_model_name_or_path=stabilityai/stable-diffusion-xl-base-1.0",
        f"--instance_data_dir={instance_data_dir}",
        f"--output_dir={b_lora_trained_folder}",
        f"--instance_prompt={instance_prompt}",
        "--resolution=1024",
        "--rank=64",
        "--train_batch_size=1",
        "--learning_rate=5e-5",
        "--lr_scheduler=constant",
        "--lr_warmup_steps=0",
        f"--max_train_steps={max_train_steps}",
        f"--checkpointing_steps={checkpoint_steps}",
        "--seed=0",
        "--gradient_checkpointing",
        "--use_8bit_adam",
        "--mixed_precision=fp16",
        "--push_to_hub",
        f"--hub_token={hf_token}"
    ]

    try:
        subprocess.run(command, check=True)
        print("Training is finished!")
    
    except subprocess.CalledProcessError as e:
        print(f"An error occurred: {e}")

def main(image_path, b_lora_trained_folder, instance_prompt):
    
    local_dir = "image_to_train"
    # Check if the directory exists and create it if necessary
    if not os.path.exists(local_dir):
        os.makedirs(local_dir)

    shutil.copy(image_path, local_dir)
    print(f"source image has been copied in {local_dir} directory")
    
    max_train_steps = 1000
    checkpoint_steps = 500
    
    train_dreambooth_blora_sdxl(local_dir, b_lora_trained_folder, instance_prompt, max_train_steps, checkpoint_steps)
    
    your_username = api.whoami(token=hf_token)["name"]
    
    return f"Done, your trained model has been stored in your models library: {your_username}/{b_lora_trained_folder}"

css = """
#col-container {max-width: 780px; margin-left: auto; margin-right: auto;}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        image = gr.Image(sources=["upload"], type="filepath")
        b_lora_name = gr.Textbox(label="b_lora_name", placeholder="b_lora_trained_folder")
        instance_prompt = gr.Textbox(label="instance prompt", placeholder="[v42]")
        train_btn = gr.Button("Train B-LoRa")
        status = gr.Textbox(label="status")
    
    train_btn.click(
        fn = main,
        inputs = [image, b_lora_name, instance_prompt],
        outputs = [status]
    )

demo.launch(debug=True)