#!/usr/bin/env python

import os
import argparse

import gradio as gr
import torch

from app_stylization import demo as demo_stylization
from app_zero_shot import demo as demo_zero_shot

DESCRIPTION = "# [BLIP-Diffusion](https://github.com/salesforce/LAVIS/tree/main/projects/blip-diffusion)"

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(
        value="Duplicate Space for private use",
        elem_id="duplicate-button",
        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
    )
    with gr.Tabs():
        with gr.Tab(label="Zero-shot Subject-driven Generation"):
            demo_zero_shot.render()
        with gr.Tab(label="Stylization"):
            demo_stylization.render()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--share", action="store_true", description="Share the app via Gradio")
    args = parser.parse_args()

    share = True if args.share else False
    demo.queue(max_size=20).launch(share=share)