#!/usr/bin/env python

import gradio as gr
import torch
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from app_image_to_3d import create_demo as create_demo_image_to_3d
from app_text_to_3d import create_demo as create_demo_text_to_3d
from model import Model

DESCRIPTION = "# [Shap-E](https://github.com/openai/shap-e)"

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

model = Model()

# Create a FastAPI app
app = FastAPI()

# Add CORS middleware to allow cross-origin requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

with gr.Blocks(css_paths="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Tabs():
        with gr.Tab(label="Text to 3D"):
            create_demo_text_to_3d(model)
        with gr.Tab(label="Image to 3D"):
            create_demo_image_to_3d(model)

# Mount the Gradio app
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    # For local development
    demo.queue(max_size=10).launch()