from diffusers import StableDiffusionXLPipeline, AutoencoderKL
import torch
import random
import os
import gradio as gr

hf_token = os.getenv("HF_TOKEN")
model_id = int(os.getenv("Model"))
nsfw_filter_enabled = int(os.getenv("Safe"))
naughty_words = os.getenv("NaughtyWords").split()
override = os.getenv("Override")

#stable-diffusion-xl-base-1.0  0 - base model
#Colossus_Project_XL           1 - better people
#AlbedoBaseXL_v11              2 - realistic
#JuggernautXL_v7               3 - better faces
#RealVisXL_V2.0                4 - better photorealism

model_url_list = ["stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors",
                 "Krebzonide/Colossus_Project_XL/blob/main/colossusProjectXLSFW_v202BakedVAE.safetensors",
                 "Krebzonide/AlbedoBaseXL_v11/blob/main/albedobaseXL_v11.safetensors",
                 "Krebzonide/JuggernautXL_version5/blob/main/juggernautXL_v7Rundiffusion.safetensors",
                 "SG161222/RealVisXL_V2.0/blob/main/RealVisXL_V2.0.safetensors",
                 "Krebzonide/AcornIsSpinning_acornXLV1/blob/main/acornIsSpinning_acornxlV1.safetensors"]

css = """
.btn-green {
  background-image: linear-gradient(to bottom right, #6dd178, #00a613) !important;
  border-color: #22c55e !important;
  color: #166534 !important;
}
.btn-green:hover {
  background-image: linear-gradient(to bottom right, #6dd178, #6dd178) !important;
}
"""

def generate(prompt, neg_prompt, samp_steps, cfg_scale, batch_size, seed, height, width, progress=gr.Progress(track_tqdm=True)):
    print("---------------------------------------------------------")
    print(prompt)
    print(neg_prompt)
    prompt = prompt.lower()
    if nsfw_filter_enabled:
        if prompt[:len(override)] == override:
            prompt = prompt[len(override):]
        else:
            neg_prompt = neg_prompt + ", child, nsfw, nude, underwear"
            for word in naughty_words:
                if prompt.find(word) >= 0:
                    return None, 58008
    if seed < 0:
        seed = random.randint(1,999999)
    images = pipe(
        prompt,
        negative_prompt=neg_prompt,
        num_inference_steps=samp_steps,
        guidance_scale=cfg_scale,
        num_images_per_prompt=batch_size,
        height=height,
        width=width,
        generator=torch.manual_seed(seed),
    ).images
    return gr.update(value = [(img, f"Image {i+1}") for i, img in enumerate(images)], height=height+90), seed
        
def set_base_model(base_model_id):
    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
    global model_url_list
    model_url = "https://huggingface.co/" + model_url_list[base_model_id]
    pipe = StableDiffusionXLPipeline.from_single_file(
        model_url,
        torch_dtype = torch.float16,
        variant = "fp16",
        vae = vae,
        use_safetensors = True,
        use_auth_token=hf_token
    )
    pipe.to("cuda")
    pipe.enable_xformers_memory_efficient_attention()
    return pipe

def update_pixel_ratio(num1, num2):
    return [round((num1-(num1%8))*num2/1048576,3), num1-(num1%8)]

examples = [
    ['A group of 4 students from University of Wisconsin Stout sitting at a table talking, men and women, detailed faces, focused',
    'glitch, deformed, cross-eyed'],
    ['A serious capybara at work, wearing a suit',
    'low quality'],
    ['a graffiti of a robot serving meals to people',
    'low quality'],
    ['photo of a small cozy modern house in red woods on a mountain, solar panels, garage, driveway, great view, sunshine',
    'red house'],
    ['cinematic photo of a woman sitting at a cafe, 35mm photograph, film, bokeh, professional, 4k, detailed face',
    'drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly'],
    ['analog film photo of old woman on the streets of london, faded film, desaturated, 35mm photo, grainy, vignette, vintage, Kodachrome, Lomography, stained, highly detailed, found footage',
    'painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured']
]

with gr.Blocks(css=css) as demo:
    with gr.Column():
        prompt = gr.Textbox(label="Prompt")
        negative_prompt = gr.Textbox(label="Negative Prompt")
        submit_btn = gr.Button("Generate", elem_classes="btn-green")
        with gr.Row():
            samp_steps = gr.Slider(1, 30, value=20, step=1, label="Sampling steps")
            cfg_scale = gr.Slider(1, 10, value=4, step=0.5, label="Guidance scale")
            batch_size = gr.Slider(1, 2, value=1, step=1, label="Batch size", interactive=True)
        with gr.Row():
            height = gr.Slider(label="Height", value=1024, minimum=8, maximum=1536, step=8)
            width = gr.Slider(label="Width", value=1024, minimum=8, maximum=1024, step=8)
        with gr.Row():
            pixels = gr.Number(label="Pixel Ratio", value=1, interactive=False)
            seed = gr.Number(label="Seed", value=-1, minimum=-1, precision=0)
        gallery = gr.Gallery(show_label=False, preview=True, container=False, height=1100)
        with gr.Row():
            lastSeed = gr.Number(label="Last Seed", value=-1, interactive=False)
        ex = gr.Examples(examples=examples, inputs=[prompt, negative_prompt])
    submit_btn.click(generate, [prompt, negative_prompt, samp_steps, cfg_scale, batch_size, seed, height, width], [gallery, lastSeed], queue=True)
    height.release(update_pixel_ratio, [height, width], [pixels, height], queue=False)
    width.release(update_pixel_ratio, [width, height], [pixels, width], queue=False)

pipe = set_base_model(model_id)
demo.launch(debug=True)