import gradio as gr
import numpy as np
import random
import torch
import spaces
import re
from diffusers import (
    DiffusionPipeline,
    AutoencoderTiny,
)
from huggingface_hub import hf_hub_download

import config
styles_name = [style["name"] for style in config.style_list]
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2240

def feifeimodload():

    dtype = torch.bfloat16
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    pipe = DiffusionPipeline.from_pretrained(
        "aifeifei798/DarkIdol-flux-v1", torch_dtype=dtype
    ).to(device)
    
    pipe.load_lora_weights(
        hf_hub_download("aifeifei798/feifei-flux-lora-v1.1", "feifei-v1.1.safetensors"),
        adapter_name="feifei",
    )
    pipe.load_lora_weights(
        hf_hub_download("aifeifei798/sldr_flux_nsfw_v2-studio", "sldr_flux_nsfw_v2-studio.safetensors"),
        adapter_name="sldr_flux_nsfw_v2",
    )
    pipe.vae.enable_slicing()
    pipe.vae.enable_tiling()
    torch.cuda.empty_cache()
    return pipe
    
pipe = feifeimodload()

@spaces.GPU()
def infer(prompt="",  styles_Radio="(None)", Expansion_select=False, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidancescale=3.5, num_feifei=0.35, nsfw_select=False, nsfw_slider=1, progress=gr.Progress(track_tqdm=True)):
    
    if nsfw_select:
        pipe.set_adapters(
            ["feifei", "sldr_flux_nsfw_v2"],
            adapter_weights=[num_feifei,nsfw_slider],
        )
        pipe.fuse_lora(
            adapter_name=["feifei", "sldr_flux_nsfw_v2"],
            lora_scale=1.0,
        )
    else:
        pipe.set_adapters(
            ["feifei"],
            adapter_weights=[num_feifei],
        )
        pipe.fuse_lora(
            adapter_name=["feifei"],
            lora_scale=1.0,
        )
        
    if Expansion_select:
        # Define the replacement string
        replacement = " feifei, A beautiful, 18 yo kpop idol, large-busted Japanese slim girl, with light makeup, gazing deeply into the camera, "
        
        # Perform the replacements with re.IGNORECASE
        prompt = re.sub(r"girl", replacement, prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"young woman", replacement, prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"woman", replacement, prompt, flags=re.IGNORECASE)
        prompt = re.sub(r"model", replacement, prompt, flags=re.IGNORECASE)
        
    if styles_Radio:
        style_name = styles_Radio
        for style in config.style_list:
            if style["name"] == style_name:
                prompt = style["prompt"].replace("{prompt}", prompt)
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    pipe.set_adapters(
        ["feifei"],
        adapter_weights=[num_feifei],
    )
    pipe.fuse_lora(
        adapter_name=["feifei"],
        lora_scale=1.0,
    )
        
    image = pipe(
            prompt = "flux, 8k, ",
            prompt_2 = prompt,
            width = width,
            height = height,
            num_inference_steps = num_inference_steps, 
            generator = generator,
            guidance_scale=guidancescale
    ).images[0] 
    return image, seed
 
examples = [
    "this photo is a girl",
    "this photo is a girl in bikini",
    "this photo is a cute girl in cute bikini",
    "girl, sunrise",
    "DarkIdol flux girl",
    "a sexy girl,poses,look at camera,Slim figure, gigantic breasts,poses,natural,High-quality photography, creative composition, fashion foresight, a strong visual style, and an aura of luxury and sophistication collectively define the distinctive aesthetic of Vogue magazine.",
    "real model slight smile girl in real life",
    "real model smile girl in real life",
    "real model girl in real life",
    "A high-resolution photograph of girl in a serene, natural setting, with soft, warm lighting, and a minimalist aesthetic, showcasing a elegant fragrance bottle and the model's effortless, emotive expression, with impeccable styling, and a muted color palette, evoking a sense of understated luxury and refinement."
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# Duke86Syl """)
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=12,
                placeholder="Enter your prompt",
                value="",
                container=False,
            )
        run_button = gr.Button("Run")
        with gr.Row():
            styles_Radio = gr.Dropdown(
                styles_name,
                label="Styles",
                multiselect=False,
                value="(None)",
            )
        result = gr.Image(label="Result", show_label=False,height=520)
        Expansion_select = gr.Checkbox(label="FeiFei Expansion") 
        with gr.Accordion("Advanced Settings", open=False):
            
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():
                gr.Markdown(''' - 21:9 2240x1024
                            - 16:9 1856x1024
                            - 9:7  1344x1024 ''')                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1344,
                )
            
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )

            with gr.Row():
                guidancescale = gr.Slider(
                    label="Guidance scale",
                    minimum=0,
                    maximum=10,
                    step=0.1,
                    value=3.5,
                )
                
            with gr.Row():
                num_feifei = gr.Slider(
                    label="FeiFei",
                    minimum=0,
                    maximum=2,
                    step=0.05,
                    value=0.45,
                )
            with gr.Row():
                nsfw_select = gr.Checkbox(label="NSFW") 
                nsfw_slider = gr.Slider(
                    label="NSFW",
                    minimum=0,
                    maximum=2,
                    step=0.05,
                    value=0.75,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples=False
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, styles_Radio, Expansion_select, seed, randomize_seed, width, height, num_inference_steps, guidancescale, num_feifei, nsfw_select, nsfw_slider],
        outputs = [result, seed]
    )

demo.launch()