import modal

image = (
    modal.Image.debian_slim()
    .apt_install("libgl1", "libglib2.0-0", "libsm6", "libxrender1", "libxext6", "ffmpeg")
    .pip_install("torch", "torchvision", "torchaudio", index_url="https://download.pytorch.org/whl/cu128")
    .pip_install_from_pyproject("./pyproject.toml")
    .pip_install("sageattention")
    .add_local_dir(local_path="../Wan2GP", remote_path="/root/Wan2GP", ignore=["*.venv"])
)

app = modal.App(name="Wan2GP", image=image)

vol = modal.Volume.from_name("wan2gp-ckpts", create_if_missing=True, environment_name="dev")
outputs = modal.Volume.from_name("wan2gp-outputs", create_if_missing=True, environment_name="dev")

with image.imports():
    import subprocess
    import os


    # @app.function(max_containers=1, gpu="L40S",
    #               scaledown_window=900,
    #               timeout=3600,
    #               volumes={"/root/Wan2GP/ckpts": vol,
    #                        "/root/Wan2GP/outputs": outputs})
    # @modal.concurrent(max_inputs=50)
    # @modal.web_server(8000, startup_timeout=2400)
    # def ui():
    #     os.chdir("/root/Wan2GP")
    #     subprocess.Popen(
    #         "python /root/Wan2GP/wgp.py --listen --server-port 8000 --compile --attention sage --profile 1",
    #         shell=True)

    @app.function(max_containers=1, gpu="L40S",
                  scaledown_window=900,
                  timeout=60 * 60 * 18,
                  volumes={"/root/Wan2GP/ckpts": vol,
                           "/root/Wan2GP/outputs": outputs})
    @modal.concurrent(max_inputs=1000)
    @modal.asgi_app()
    def serve():
        import asyncio
        import atexit
        from fastapi import FastAPI
        from gradio import mount_gradio_app
        os.chdir("/root/Wan2GP")
        import sys
        sys.path.append('/root/Wan2GP')
        sys.argv = ["--listen", "--server-port", "8000", "--compile", "--attention", "sage", "--profile", "1"]
        from wgp import autosave_queue, create_ui, save_path

        app = FastAPI()

        atexit.register(autosave_queue)
        # download_ffmpeg()
        # threading.Thread(target=runner, daemon=True).start()
        os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
        server_port = 8000
        if os.name == "nt":
            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
        server_name = "0.0.0.0"
        demo = create_ui()
        demo.queue(max_size=5)
        return mount_gradio_app(app=app, blocks=demo, path="/",
                                server_name=server_name, server_port=server_port,
                                allowed_paths=[save_path])