from PIL import Image
import gradio as gr
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    AutoImageProcessor,
    AutoModel,
)
from transformers.generation.configuration_utils import GenerationConfig
from transformers.generation import (
    LogitsProcessorList,
    PrefixConstrainedLogitsProcessor,
    UnbatchedClassifierFreeGuidanceLogitsProcessor,
)
import torch
from emu3.mllm.processing_emu3 import Emu3Processor
import spaces

# Model paths
EMU_GEN_HUB = "BAAI/Emu3-Gen"
EMU_CHAT_HUB = "BAAI/Emu3-Chat"
VQ_HUB = "BAAI/Emu3-VisionTokenizer"

# Prepare models and processors
# Emu3-Gen model and processor
gen_model = AutoModelForCausalLM.from_pretrained(
    EMU_GEN_HUB,
    device_map="cuda:0",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    trust_remote_code=True,
)

gen_tokenizer = AutoTokenizer.from_pretrained(EMU_GEN_HUB, trust_remote_code=True)
gen_image_processor = AutoImageProcessor.from_pretrained(
    VQ_HUB, trust_remote_code=True
)
gen_image_tokenizer = AutoModel.from_pretrained(
    VQ_HUB, device_map="cuda:0", trust_remote_code=True
).eval()
gen_processor = Emu3Processor(gen_image_processor, gen_image_tokenizer, gen_tokenizer)

# Emu3-Chat model and processor
chat_model = AutoModelForCausalLM.from_pretrained(
    EMU_CHAT_HUB,
    device_map="cuda:0",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    trust_remote_code=True,
)

chat_tokenizer = AutoTokenizer.from_pretrained(EMU_CHAT_HUB, trust_remote_code=True)
chat_image_processor = AutoImageProcessor.from_pretrained(
    VQ_HUB, trust_remote_code=True
)
chat_image_tokenizer = AutoModel.from_pretrained(
    VQ_HUB, device_map="cuda:0", trust_remote_code=True
).eval()
chat_processor = Emu3Processor(
    chat_image_processor, chat_image_tokenizer, chat_tokenizer
)

@spaces.GPU(duration=120)
def generate_image(prompt):
    POSITIVE_PROMPT = " masterpiece, film grained, best quality."
    NEGATIVE_PROMPT = (
        "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, "
        "fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, "
        "signature, watermark, username, blurry."
    )

    classifier_free_guidance = 3.0
    full_prompt = prompt + POSITIVE_PROMPT

    kwargs = dict(
        mode="G",
        ratio="1:1",
        image_area=gen_model.config.image_area,
        return_tensors="pt",
    )
    pos_inputs = gen_processor(text=full_prompt, **kwargs)
    neg_inputs = gen_processor(text=NEGATIVE_PROMPT, **kwargs)

    # Prepare hyperparameters
    GENERATION_CONFIG = GenerationConfig(
        use_cache=True,
        eos_token_id=gen_model.config.eos_token_id,
        pad_token_id=gen_model.config.pad_token_id,
        max_new_tokens=40960,
        do_sample=True,
        top_k=2048,
    )

    h, w = pos_inputs.image_size[0]
    constrained_fn = gen_processor.build_prefix_constrained_fn(h, w)
    logits_processor = LogitsProcessorList(
        [
            UnbatchedClassifierFreeGuidanceLogitsProcessor(
                classifier_free_guidance,
                gen_model,
                unconditional_ids=neg_inputs.input_ids.to("cuda:0"),
            ),
            PrefixConstrainedLogitsProcessor(
                constrained_fn,
                num_beams=1,
            ),
        ]
    )

    # Generate
    outputs = gen_model.generate(
        pos_inputs.input_ids.to("cuda:0"),
        generation_config=GENERATION_CONFIG,
        logits_processor=logits_processor,
    )

    mm_list = gen_processor.decode(outputs[0])
    for idx, im in enumerate(mm_list):
        if isinstance(im, Image.Image):
            return im
    return None

def vision_language_understanding(image, text):
    inputs = chat_processor(
        text=text,
        image=image,
        mode="U",
        padding_side="left",
        padding="longest",
        return_tensors="pt",
    )

    # Prepare hyperparameters
    GENERATION_CONFIG = GenerationConfig(
        pad_token_id=chat_tokenizer.pad_token_id,
        bos_token_id=chat_tokenizer.bos_token_id,
        eos_token_id=chat_tokenizer.eos_token_id,
        max_new_tokens=320,
    )

    # Generate
    outputs = chat_model.generate(
        inputs.input_ids.to("cuda:0"),
        generation_config=GENERATION_CONFIG,
        max_new_tokens=320,
    )

    outputs = outputs[:, inputs.input_ids.shape[-1] :]
    response = chat_processor.batch_decode(outputs, skip_special_tokens=True)[0]
    return response

def chat(history, user_input, user_image):
    if user_image is not None:
        # Use Emu3-Chat for vision-language understanding
        response = vision_language_understanding(user_image, user_input)
        # Append the user input and response to the history
        history = history + [(user_input, response)]
    else:
        # Use Emu3-Gen for image generation
        generated_image = generate_image(user_input)
        if generated_image is not None:
            # Append the user input and generated image to the history
            history = history + [(user_input, generated_image)]
        else:
            # If image generation failed, respond with an error message
            history = history + [
                (user_input, "Sorry, I could not generate an image.")
            ]
    return history, history, gr.update(value=None)

def clear_input():
    return gr.update(value="")

with gr.Blocks() as demo:
    gr.Markdown("# Emu3 Chatbot Demo")
    gr.Markdown(
        "This is a chatbot demo for image generation and vision-language understanding using Emu3 models."
    )

    chatbot = gr.Chatbot()
    state = gr.State([])
    with gr.Row():
        with gr.Column(scale=0.85):
            user_input = gr.Textbox(
                show_label=False, placeholder="Type your message here...", lines=2
            ).style(container=False)
        with gr.Column(scale=0.15, min_width=0):
            submit_btn = gr.Button("Send")
    user_image = gr.Image(
        source="upload", type="pil", label="Upload an image (optional)"
    )

    submit_btn.click(
        chat,
        inputs=[state, user_input, user_image],
        outputs=[chatbot, state, user_image],
    ).then(fn=clear_input, inputs=[], outputs=user_input)
    user_input.submit(
        chat,
        inputs=[state, user_input, user_image],
        outputs=[chatbot, state, user_image],
    ).then(fn=clear_input, inputs=[], outputs=user_input)

demo.launch()