# Ref: https://huggingface.co/spaces/multimodalart/cosxl
import gradio as gr
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
import spaces 
import torch 
import os

from compel import Compel, ReturnedEmbeddingsType

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

model_id = "aipicasso/emi-2"
token=os.environ["TOKEN"]

scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_id,subfolder="scheduler",token=token)
pipe_normal = StableDiffusionXLPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.bfloat16,token=token)

negative_ti_file = hf_hub_download(repo_id="Aikimi/unaestheticXL_Negative_TI", filename="unaestheticXLv31.safetensors")
state_dict = load_file(negative_ti_file)
pipe_normal.load_textual_inversion(state_dict["clip_g"], token="unaestheticXLv31", text_encoder=pipe_normal.text_encoder_2, tokenizer=pipe_normal.tokenizer_2)
pipe_normal.load_textual_inversion(state_dict["clip_l"], token="unaestheticXLv31", text_encoder=pipe_normal.text_encoder, tokenizer=pipe_normal.tokenizer)

pipe_normal.to("cuda")

pipe_normal.enable_freeu(s1=1.2, s2=0.7, b1=1.1, b2=1.3)

compel = Compel(tokenizer=[pipe_normal.tokenizer, pipe_normal.tokenizer_2] , 
                text_encoder=[pipe_normal.text_encoder, pipe_normal.text_encoder_2], 
                returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, 
                requires_pooled=[False, True])

@spaces.GPU
def run_normal(prompt, negative_prompt="", guidance_scale=7.5, progress=gr.Progress(track_tqdm=True)):
    if(prompt==""):
        conditioning, pooled = compel("1girl, (upper body)++, brown bob short hair, brown eyes, looking at viewer, cherry blossom")
    else:
        conditioning, pooled = compel(prompt)
    negative_conditioning, negatice_pooled = compel("unaestheticXLv31, bad hand, bad anatomy, low quality, 3d, photo, realism, text, sign, "+negative_prompt)
    result = pipe_normal(
        prompt_embeds=conditioning,
        pooled_prompt_embeds=pooled, 
        negative_prompt_embeds=negative_conditioning, 
        negative_pooled_prompt_embeds=negatice_pooled,
        num_inference_steps = 25,
        guidance_scale = guidance_scale,
        width = 768,
        height = 1344)
    
    return result.images[0]

css = '''
.gradio-container{
max-width: 768px !important;
margin: 0 auto;
}
'''

normal_examples = [
    "1girl, (upper body)++, brown bob short hair, brown eyes, looking at viewer, cherry blossom",
    "1girl, (full body)++, brown bob short hair, brown eyes, school uniform, cherry blossom",
    "no humans, manga, black and white, monochrome, Mt. fuji, 4k, highly detailed",
    "no humans, manga, black and white, monochrome, Shibuya street, 4k, highly detailed",
    "1boy, (upper body)++, black short hair, black eyes, looking at viewer, green leaves",
    "1boy, (full body)++, black bob short hair, black eyes, school uniform, green leaves",
]

with gr.Blocks(css=css) as demo:
    gr.Markdown('''# Emi 2
    Official demo for [Emi 2](https://huggingface.co/aipicasso/emi-2). Click the generate button!
    ''')
    with gr.Group():
        with gr.Row():
          prompt_normal = gr.Textbox(show_label=False, scale=4, placeholder="Your prompt, e.g.: 1girl, (upper body)++, brown bob short hair, brown eyes, looking at viewer, cherry blossom")
          button_normal = gr.Button("Generate", min_width=120)
        output_normal = gr.Image(label="Your result image", interactive=False)
        with gr.Accordion("Advanced Settings", open=False):
          negative_prompt_normal = gr.Textbox(label="Negative Prompt")
          guidance_scale_normal = gr.Number(label="Guidance Scale", value=7.5)
    gr.Examples(examples=normal_examples, fn=run_normal, inputs=[prompt_normal], outputs=[output_normal], cache_examples=True) 
    
    gr.on(
        triggers=[
            button_normal.click,
            prompt_normal.submit
        ],
        fn=run_normal,
        inputs=[prompt_normal, negative_prompt_normal, guidance_scale_normal],
        outputs=[output_normal],
    )
if __name__ == "__main__":
    demo.launch(share=True)