# 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 torch
from diffusers import Transformer2DModel, PixArtSigmaPipeline, AutoencoderKL, DPMSolverMultistepScheduler, DDIMScheduler, EulerAncestralDiscreteScheduler, DPMSolverSDEScheduler
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, QuantoConfig, EetqConfig

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

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 = [
    "芝生の上にあるピザ",
    "東京の桜と建物。満開の桜の木が並び、ピンク色の花びらが風に舞っている。桜の背景には東京の高層ビルや伝統的な建物が調和して立っている。春の陽光が全体を明るく照らし、桜と建物が美しく映えている。都市の活気と自然の美しさが融合した風景。",
    "パリは燃えているか",
]

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