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



# Set up the argument parser
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()

# Initialize lists for interfaces and tab names
interface_list, tab_names = [], []

# Conditional blocks to add chat and generation interfaces
if not args.disable_chat:
    demo_chat = build_chat(args)
    interface_list.append(demo_chat)
    tab_names.append("Multimodal Chat")

if not args.disable_generate:
    demo_generation = build_generation(args)
    interface_list.append(demo_generation)
    tab_names.append("Multimodal Generation")

# Create a tabbed interface with the specified interfaces and tabs
demo_all = gr.TabbedInterface(
    interface_list=interface_list,
    tab_names=tab_names,
    title=args.title,
    theme=gr.themes.Default(primary_hue="blue", secondary_hue="blue"),
)

# Launch the application with the specified settings
demo_all.queue(
    concurrency_count=args.concurrency_count,
    status_update_rate=3,
    api_open=False,
).launch(
    enable_queue=True,
    share=args.share,
)