# 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=0.9, s2=0.2, b1=1.3, b2=1.4)

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)):
    conditioning, pooled = compel(prompt)
    negative_conditioning, negatice_pooled = compel("unaestheticXLv31-, bad hand, bad anatomy, low quality, "+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 person, manga, black and white, monochrome, Mt. fuji, 4k, highly, detailed",
    "no person, manga, black and white, monochrome, buldings in Tokyo from sky, 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
    ''')
    with gr.Group():
        with gr.Row():
          prompt_normal = gr.Textbox(show_label=False, scale=4, placeholder="Your prompt, e.g.: 1girl, face, brown bob short hair, brown eyes, looking at viewer")
          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)