# 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")

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, "unaestheticXLv31--, "+neg_prompt])
    
    result = pipe(
        prompt_embeds=conditioning[0:1],
        pooled_prompt_embeds=pooled[0:1], 
        negative_prompt_embeds=conditioning[1:2], 
        negative_pooled_prompt_embeds=pooled[1:2],
        num_inference_steps = 20,
        guidance_scale = guidance_scale,
        width = 1344,
        height = 768)
    
    return result.images[0]

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

normal_examples = ["1girl, face, brown bob short hair, brown eyes, looking at viewer"]
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)
    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)