#!/usr/bin/env python

import gradio as gr
import torch

torch.jit.script = lambda x: x

from app_stylization import demo as demo_stylization  # noqa: E402
from app_zero_shot import demo as demo_zero_shot  # noqa: E402

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_paths="style.css") as demo:
    gr.Markdown(DESCRIPTION)

    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__":
    demo.queue(max_size=20).launch()