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

model_id = int(os.getenv("Model"))
nsfw_filter = int(os.getenv("Safe"))

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

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_version5.safetensors",
                 "Krebzonide/JuggernautXL_version5/blob/main/juggernautXL_version6Rundiffusion.safetensors",
                 "SG161222/RealVisXL_V2.0/blob/main/RealVisXL_V2.0.safetensors",
                 "Krebzonide/AcornIsSpinning_acornXLV1/blob/main/acornIsSpinning_acornxlV1.safetensors"]

naughtyWords = ["nude", "nsfw", "naked", "porn", "boob", "tit", "nipple", "vagina", "pussy", "panties", "underwear", "upskirt", "bottomless", "topless", "petite", "xxx"]


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)):
    prompt = prompt.lower()
    if nsfw_filter:
        if prompt[:10] == "krebzonide":
            prompt = prompt[10:]
        else:
            neg_prompt = neg_prompt + ", child, nsfw, nipples, nude, underwear"
            for word in naughtyWords:
                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 [(img, f"Image {i+1}") for i, img in enumerate(images)], 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="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"
    )
    pipe.to("cuda")
    return pipe

def update_pixel_ratio(num1, num2):
    return [round(num1*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']#,
    #['nude photo of a 20 year old model in the back seat of a car, detailed face',
    #'big boobs'],
    #['nude photo of a 20 year old man, penis and testicles, dick and balls, erection',
    #'woman']
]

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, 6, value=1, step=1, label="Batch size", interactive=False)
        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(label="Generated images")
        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)