# -*- coding: utf-8 -*-

# ===================================================
#
#    Author        : Fan Zhang
#    Email         : zhangfan@baai.ac.cn
#    Institute     : Beijing Academy of Artificial Intelligence (BAAI)
#    Create On     : 2023-12-11 15:34
#    Last Modified : 2024-01-04 10:15
#    File Name     : app.py
#    Description   :
#
# ===================================================

import os
import sys
import logging

# 设置详细日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

os.system("pip uninstall -y gradio")
os.system("pip install gradio==6.5.1")

import argparse
import gradio as gr
from demo.generation_frontend import build_generation
from demo.chat_frontend import build_chat

parser = argparse.ArgumentParser()
parser.add_argument("--title", type=str, default='Emu')
parser.add_argument("--host", type=str, default="0.0.0.0")
parser.add_argument("--port", type=int, default=9002)
parser.add_argument("--share", action="store_true")
parser.add_argument("--controller-url", type=str, default="http://218.91.113.230:9003")
parser.add_argument("--concurrency-count", type=int, default=8)
parser.add_argument("--disable-chat", action="store_true")
parser.add_argument("--disable-generate", action="store_true")
args = parser.parse_args()

if __name__ == "__main__":
    logger.info("Starting Emu2 application...")
    
    title = "Emu2: Generative Multimodal Models are In-Context Learners<br> \
            <h2 align='center'> \
            [<a href='https://baaivision.github.io/emu2' target='_blank' rel='noopener'>project page</a>] \
            [<a href='https://github.com/baaivision/Emu' target='_blank' rel='noopener'>code</a>] \
            [<a href='https://arxiv.org/abs/2312.13286' target='_blank' rel='noopener'>paper</a>] \
            </h2> \
            <div align='center'> \
            <font size=4>Experience Emu2 with 👉</font> \
            <font size=6>[<b><a href='https://emu.ssi.plus' target='_blank' rel='noopener'>FAST demo</a>]</b></font> \
            <font size=4>👈 <b>WAY MORE FASTER!!!</b></font> \
            </div> \
            <h3 align='center'> \
            🤗HF models: \
            <a href='https://huggingface.co/BAAI/Emu2' target='_blank' rel='noopener'>Emu2</a> | \
            <a href='https://huggingface.co/BAAI/Emu2-Chat' target='_blank' rel='noopener'>Emu2-Chat</a> | \
            <a href='https://huggingface.co/BAAI/Emu2-Gen' target='_blank' rel='noopener'>Emu2-Gen</a> \
            </h3> \
            <h4 align='center'> \
            [<a href='https://jwolpxeehx.feishu.cn/docx/KskPdU99FomufKx4G9hcQMeQnHv' target='_blank' rel='noopener'>使用说明</a>] \
            [<a href='https://jwolpxeehx.feishu.cn/docx/RYHNd1tvEo8k8Mx9HeMcvvxWnvZ' target='_blank' rel='noopener'>User Guide</a>] \
            </h4>"

    interface_list, tab_names = [], []
    
    try:
        if not args.disable_chat:
            logger.info("Building chat interface...")
            demo_chat = build_chat(args)
            interface_list.append(demo_chat)
            tab_names.append("Multimodal Chat")
        
        if not args.disable_generate:
            logger.info("Building generation interface...")
            demo_generation = build_generation(args)
            interface_list.append(demo_generation)
            tab_names.append("Multimodal Generation")
        
        if not interface_list:
            logger.error("No interfaces enabled. Enable at least one of --disable-chat or --disable-generate.")
            sys.exit(1)
            
        logger.info(f"Creating TabbedInterface with {len(interface_list)} tabs")
        demo_all = gr.TabbedInterface(
            interface_list=interface_list,
            tab_names=tab_names,
            title=title,
        )
        
        logger.info(f"Launching application on {args.host}:{args.port}")
        demo_all.launch(
            server_name=args.host,
            server_port=args.port,
            share=args.share,
            theme=gr.themes.Default(primary_hue="blue", secondary_hue="blue"),
            ssr_mode=False,  # 禁用 SSR
            debug=True,      # 启用调试
        )
        
    except Exception as e:
        logger.error(f"Application failed to start: {e}", exc_info=True)
        sys.exit(1)