import subprocess
import sys

from DepthFlow.Webui import ESTIMATORS, UPSCALERS, DepthGradio

from Broken import Runtime
from Broken.Externals.Depthmap import DepthAnythingV2

# -----------------------------------------------|

if Runtime.ZeroGPU:

    # Install latest git main code
    for dependency in (
        "git+https://github.com/BrokenSource/BrokenSource@main",
        "git+https://github.com/BrokenSource/ShaderFlow@main",
        "git+https://github.com/BrokenSource/DepthFlow@main",
    ):
        subprocess.run((
            sys.executable, "-m", "uv", "pip", "install",
            dependency, "--upgrade",
        ))

    # Remove legacy upscalers
    UPSCALERS.pop("Real-ESRGAN")
    UPSCALERS.pop("Waifu2x")

    # Remove unsuitable models
    ESTIMATORS.clear()
    ESTIMATORS.update({
        # Base and Large models are CC BY-NC 4.0, better safe than sorry
        "DepthAnything2": DepthAnythingV2(model=DepthAnythingV2.Model.Small),
    })

# -----------------------------------------------|

DepthGradio().launch(
    workers=4,
    nvenc=True,
    turbo=True,
    pwa=True,
)