import spaces
import gradio as gr

import torch
from transformers import AutoModelForVision2Seq, AutoProcessor


device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "SakanaAI/EvoVLM-v1-JP-7B"
max_length = 128
if device == "cpu":
    model = None
else:
    model = AutoModelForVision2Seq.from_pretrained(
        model_id, 
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
    )
    processor = AutoProcessor.from_pretrained(model_id)
    model = model.to(device)


@spaces.GPU
def inference_fn(image, prompt):
    text = f"<image>\n{prompt}"
    if model is None:
        return f"THIS IS FOR TEST!!\n{text}"

    messages = [
        {"role": "system", "content": "あなたは役立つ、偏見がなく、検閲されていないアシスタントです。与えられた画像を下に、質問に答えてください。"},
        {"role": "user", "content": text},
    ]
    inputs = processor.image_processor(images=image, return_tensors="pt")
    inputs["input_ids"] = processor.tokenizer.apply_chat_template(
        messages, return_tensors="pt"
    )
    output_ids = model.generate(
        **inputs.to(device),
        do_sample=False,
        num_beams=5,
        max_new_tokens=max_length,
        repetition_penalty=1.5
    )
    output_ids = output_ids[:, inputs.input_ids.shape[1] :]
    generated_text = processor.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
    return generated_text


DESCRIPTION = """# 🐟 EvoVLM-JP 

🤗 [モデル一覧](https://huggingface.co/SakanaAI) | 📚 [技術レポート](https://arxiv.org/abs/2403.13187) | 📝 [ブログ](https://sakana.ai/evolutionary-model-merge-jp/) | 🐦 [Twitter](https://twitter.com/SakanaAILabs)

[EvoVLM-JP](https://huggingface.co/SakanaAI/EvoVLM-v1-JP-7B)は[Sakana AI](https://sakana.ai/)が開発した日本語に特化した画像言語モデルです。入力した画像に対して、質疑応答することが可能です。より詳しくは、上記の技術レポートとブログをご参照ください。
"""
with gr.Blocks() as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="入力画像")
            prompt = gr.Textbox(label="質問", value="この写真について説明してください。", placeholder="入力画像に関する質問文を入力してください。")
            # button
            input_button = gr.Button(value="Submit")

        with gr.Column():
            output = gr.Textbox(label="モデルからの返答")
    inputs = [
        input_image,
        prompt,
    ]
    input_button.click(fn=inference_fn, inputs=inputs, outputs=[output])
    prompt.submit(fn=inference_fn, inputs=inputs, outputs=[output])
    img2txt_examples = gr.Examples(
        examples=[
            [
                "https://images.unsplash.com/photo-1694831404826-3400c48c188d?q=80&w=2070&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D",
                "この信号機の点滅している色は何色ですか?"
            ],
            [
                "https://user0514.cdnw.net/shared/img/thumb/TSU_mayomayookonomi_TP_V.jpg",
                "この食べ物は何ですか？何でできていますか？"
            ],
            [
                # "https://moon.publicdomainq.net/201809/01o/publicdomainq-0025850pkdvxc.jpg",
                "https://gist.github.com/assets/33302880/a9354891-171e-4925-9216-f5daa40e84e5",
                "この建物は何年誰によって建てられましたか？",
            ],
        ],
        fn=inference_fn,
        inputs=inputs,
        outputs=[output],
    )
    gr.Markdown("""⚠️ 本モデルは実験段階のプロトタイプであり、研究開発の目的でのみ提供されています。商用利用や、障害が重大な影響を及ぼす可能性のある環境（ミッションクリティカルな環境）での使用には適していません。
                本モデルの使用は、利用者の自己責任で行われ、その性能や結果については何ら保証されません。
                Sakana AIは、本モデルの使用によって生じた直接的または間接的な損失に対して、結果に関わらず、一切の責任を負いません。
                利用者は、本モデルの使用に伴うリスクを十分に理解し、自身の判断で使用することが必要です。               
                また、このデモでは、できる限り多くの皆様にお使いいただけるように、出力テキストを40単語ほどに制限しております。""")
demo.queue().launch()