import os
import gradio as gr
import torch
import uuid
import peft
from PIL import Image
from diffusers import AutoPipelineForText2Image, StableDiffusionXLInpaintPipeline, StableDiffusionXLPipeline
from peft import PeftModel, PeftConfig

# Define global variables
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
lora_models = {}
trigger_word = {}

# Load the pretrained model and add LoRAs
pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
pipe.to("cuda")

base_model = pipe.model
peft_config = PeftConfig.from_pretrained('lora_weights/qwe_cat_long.safetensors')
peft_model = PeftModel.from_pretrained(base_model, peft_config)
pipe.model = peft_model

# Create a dictionary of available LoRAs and their corresponding trigger words
for i in os.scandir('lora_weights'):
    if i.name != '.gitignore':
        lora_models[i.name] = i.path
        trigger_word[i.name] = i.name.split('_')[0] + ' cat bright white fur'

# Define helper functions
def save_img(image_list, prompt):
    results_folder = 'results/'
    os.makedirs(results_folder, exist_ok=True)
    for image in image_list:
        image = Image.open(image[0])
        unique_id = uuid.uuid4()
        image.save(f"{results_folder}{unique_id}.jpg")
        new_filename = f"{results_folder}{unique_id}.txt"
        with open(new_filename, "w") as file:
            file.write(prompt)

def set_lora_model(lora_name, lora_scale):
    pipe.unfuse_lora(True)
    pipe.unload_lora_weights()
    print(lora_models[lora_name])
    peft_config = PeftConfig.from_pretrained(lora_models[lora_name])
    peft_config.lora_scale = lora_scale
    peft_model = PeftModel.from_pretrained(base_model, peft_config)
    pipe.model = peft_model
    pipe.fuse_lora()
    print('Model swapped')
    return trigger_word[lora_name]

# ...

if __name__ == "__main__":
    main()