# Reference: https://huggingface.co/spaces/black-forest-labs/FLUX.1-schnell/blob/main/app.py
import spaces
import gradio as gr
import numpy as np
import random
import torch
import os
from diffusers import Transformer2DModel, PixArtSigmaPipeline, AutoencoderKL, DPMSolverMultistepScheduler, DDIMScheduler, EulerAncestralDiscreteScheduler, DPMSolverSDEScheduler
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, QuantoConfig, EetqConfig

device = "cuda"
weight_dtype = torch.bfloat16
weight_dtype_te = torch.bfloat16
MAX_SEED = np.iinfo(np.int32).max

transformer = Transformer2DModel.from_pretrained(
    "aipicasso/commonart-beta",
    torch_dtype=weight_dtype,
    token=os.environ["TOKEN"]
)

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=weight_dtype)
scheduler=DPMSolverMultistepScheduler()
pipe = PixArtSigmaPipeline(
    vae=vae,
    tokenizer=None,
    text_encoder=None,
    transformer=transformer,
    scheduler=scheduler
)

pipe.to(device)

tokenizer = AutoTokenizer.from_pretrained("cyberagent/calm2-7b")
text_encoder =  AutoModelForCausalLM.from_pretrained(
    "cyberagent/calm2-7b",
    torch_dtype=weight_dtype_te,
    device_map=device
)

@spaces.GPU()
def infer(prompt, seed=42, randomize_seed=False, width=512, height=512, num_inference_steps=20, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    with torch.no_grad():
        pos_ids = tokenizer(
            prompt, max_length=512, padding="max_length", truncation=True, return_tensors="pt",
        ).to(device)
        pos_emb = text_encoder(pos_ids.input_ids, output_hidden_states=True, attention_mask=pos_ids.attention_mask)
        pos_emb = pos_emb.hidden_states[-1]
        neg_ids = tokenizer(
            "", max_length=512, padding="max_length", truncation=True, return_tensors="pt",
        ).to(device)
        neg_emb = text_encoder(neg_ids.input_ids, output_hidden_states=True, attention_mask=neg_ids.attention_mask)
        neg_emb = neg_emb.hidden_states[-1]
    
        image = pipe(
            negative_prompt=None,
            prompt_embeds=pos_emb,
            negative_prompt_embeds=neg_emb,
            prompt_attention_mask=pos_ids.attention_mask,
            negative_prompt_attention_mask=neg_ids.attention_mask,
            max_sequence_length=512,
            width=width,
            height=height,
            num_inference_steps=num_inference_steps,
            generator=generator,
            guidance_scale=4.5).images[0]
    return image, seed
 
examples = [
    "芝生の上で燃えている大きなピザ",
    "カラフルなお花畑。赤、青、黄色、ピンクなどの花が見えて、緑色の草も見えている",
    "東京の桜と建物。満開の桜の木が並び、ピンク色の花びらが風に舞っている。桜の背景には東京の高層ビルや伝統的な建物が調和して立っている。春の陽光が全体を明るく照らし、桜と建物が美しく映えている。都市の活気と自然の美しさが融合した風景。",
    "silverでできている男性の顔写真。背景は黒い。silverはキラキラしている。",
    "ポケモンのピカチュウ。任天堂が作ったキャラクター。黄色い。ほっぺは赤い。電気を放つことができる。",
    "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."
]

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"""# CommonArt β
商用利用できる透明性の高い日本語画像生成AI [モデル](https://huggingface.co/aipicasso/commonart-beta) [ブログ](https://note.com/aipicasso/n/nf17f876839b2)
※プロンプトが短すぎると出力が乱れることがあります。
        """)
        
        with gr.Row():
            
            prompt = gr.Text(
                label="テキスト",
                show_label=False,
                max_lines=1,
                placeholder="生成したいものを日本語や英語で説明してください",
                container=False,
            )
            
            run_button = gr.Button("生成", scale=0)
        
        result = gr.Image(label="生成結果", show_label=False)
        
        with gr.Accordion("詳細設定", open=False):
            
            seed = gr.Slider(
                label="シード値",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="ランダム", value=True)
            
            with gr.Row():
                
                width = gr.Slider(
                    label="幅",
                    minimum=256,
                    maximum=768,
                    step=64,
                    value=512,
                )
                
                height = gr.Slider(
                    label="高さ",
                    minimum=256,
                    maximum=768,
                    step=64,
                    value=512,
                )
            
            with gr.Row():
                
  
                num_inference_steps = gr.Slider(
                    label="推論回数",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=20,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples="lazy"
        )

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

demo.launch()