from diffusers import StableDiffusionPipeline
from lora_diffusion import monkeypatch_lora, tune_lora_scale
import torch
import os 
import gradio as gr
import subprocess

MODEL_NAME="stabilityai/stable-diffusion-2-1-base"
INSTANCE_DIR="./data_example"
OUTPUT_DIR="./output_example"

model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
#prompt = "style of sks, baby lion"
torch.manual_seed(1)
#image = pipe(prompt, num_inference_steps=50, guidance_scale= 7).images[0] #no need
#image  # nice. diffusers are cool. #no need
#finetuned_lora_weights = "./lora_weight.pt"

#global var
counter = 0

#Getting Lora fine-tuned weights
def monkeypatching(alpha, in_prompt): #, prompt, pipe): finetuned_lora_weights
    print("****** inside monkeypatching *******")
    print(f"in_prompt is - {str(in_prompt)}")
    global counter
    if counter == 0 :
        monkeypatch_lora(pipe.unet, torch.load("./output_example/lora_weight.pt")) #finetuned_lora_weights
        tune_lora_scale(pipe.unet, alpha) #1.00)
        counter +=1
    else :
        tune_lora_scale(pipe.unet, alpha) #1.00)
    prompt = "style of hclu, " + str(in_prompt)   #"baby lion"
    image = pipe(prompt, num_inference_steps=50, guidance_scale=7).images[0]
    image.save("./illust_lora.jpg") #"./contents/illust_lora.jpg")
    return image

def accelerate_train_lora(steps):
    print("*********** inside accelerate_train_lora ***********")
    #subprocess.run(accelerate launch {"./train_lora_dreambooth.py"} \
    #subprocess.Popen(f'accelerate launch {"./train_lora_dreambooth.py"} \
    os.system( f'accelerate launch {"./train_lora_dreambooth.py"} \
      --pretrained_model_name_or_path={MODEL_NAME}  \
      --instance_data_dir={INSTANCE_DIR} \
      --output_dir={OUTPUT_DIR} \
      --instance_prompt="style of hclu" \
      --resolution=512 \
      --train_batch_size=1 \
      --gradient_accumulation_steps=1 \
      --learning_rate=1e-4 \
      --lr_scheduler="constant" \
      --lr_warmup_steps=0 \
      --max_train_steps={int(steps)}') #,shell=True) #30000
    print("*********** completing accelerate_train_lora ***********")
    return "./output_example/lora_weight.pt"
    
with gr.Blocks() as demo:
    gr.Markdown("""<h1><center>LORA - Low-rank Adaptation for Fast Text-to-Image Diffusion Fine-tuning</center></h1>
    """)
    gr.Markdown(
        """**Main Features**<br>- Fine-tune Stable diffusion models twice as faster than dreambooth method, by Low-rank Adaptation.<br>- Get insanely small end result, easy to share and download.<br>- Easy to use, compatible with diffusers.<br>- Sometimes even better performance than full fine-tuning<br>Please refer the Github repo this Space is based on, here - <a href = "https://github.com/cloneofsimo/lora">LORA</a><br>You can also refer this tweet by AK over here to quote/retweet/like, here on <a href="https://twitter.com/_akhaliq/status/1601120767009513472">Twitter</a>.<br>This Gradio Space is an attempt to explore this novel LORA approach to fine-tune Stable diffusion models, using the power and flexibility of Gradio!<br><b>To use this Space well:</b>- First, upload your set of images (4-5), then enter the number of fine-tuning steps, and then press the 'Train LORA model' button.<br>- Enter a prompt, then set the alpha value using the Slider (nearer to 1 implies overfitting to the uploaded images), and then press the 'Inference' button.<br><b>Bonus:</b>Download your fine-tuned model weights from the Gradio file component. The smaller size of LORA models (around 3-4 mb files) is the main highlight of this 'Low-rank Adaptation' approach of fine-tuning.""")

    with gr.Row():
        in_images =  gr.File(label="Upload images to fine-tune for LORA", file_count="multiple")
        #in_prompt = gr.Textbox(label="Enter a ")
        in_steps = gr.Number(label="Enter number of steps")
        in_alpha = gr.Slider(0.1,1.0, step=0.01, label="Set Alpha level - higher value has more chances to overfit")
    with gr.Row():
        b1 = gr.Button(value="Train LORA model")
        b2 = gr.Button(value="Inference using LORA model")
    with gr.Row():
        in_prompt = gr.Textbox(label="Enter a prompt for fine-tuned LORA model", visible=True)
        out_image = gr.Image(label="Image generated by LORA model")
        out_file = gr.File(label="Lora trained model weights", )
    b1.click(fn = accelerate_train_lora, inputs=in_steps, outputs=out_file)
    b2.click(fn = monkeypatching, inputs=[in_alpha, in_prompt], outputs=out_image)

demo.queue(concurrency_count=3) 
demo.launch(debug=True, show_error=True)