import os
if os.environ.get("SPACES_ZERO_GPU") is not None:
    import spaces
else:
    class spaces:
        @staticmethod
        def GPU(func):
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs)
            return wrapper
import subprocess
import gc
from pathlib import Path
from utils import set_token, get_download_file

@spaces.GPU
def fake_gpu():
  pass

HF_TOKEN = os.environ.get("HF_TOKEN")
CIVITAI_API_KEY = os.environ.get("CIVITAI_API_KEY")
set_token(HF_TOKEN)
BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = "data"
os.makedirs(str(Path(BASE_DIR, DATA_DIR)), exist_ok=True)

def get_file(url: str, path: str):
    print(f"Downloading {url} to {path}...")
    get_download_file(path, url, CIVITAI_API_KEY)

def git_clone(url: str, path: str, pip: bool=False, addcmd: str=""):
    os.makedirs(str(Path(BASE_DIR, path)), exist_ok=True)
    os.chdir(Path(BASE_DIR, path))
    print(f"Cloning {url} to {path}...")
    cmd = f'git clone {url}'
    print(f'Running {cmd} at {Path.cwd()}')
    i = subprocess.run(cmd, shell=True).returncode
    if i != 0: print(f'Error occured at running {cmd}')
    p = url.split("/")[-1]
    if not Path(p).exists: return
    if pip:
        os.chdir(Path(BASE_DIR, path, p))
        cmd = f'pip install -r requirements.txt'
        print(f'Running {cmd} at {Path.cwd()}')
        i = subprocess.run(cmd, shell=True).returncode
        if i != 0: print(f'Error occured at running {cmd}')
    if addcmd:
        os.chdir(Path(BASE_DIR, path, p))
        cmd = addcmd
        print(f'Running {cmd} at {Path.cwd()}')
        i = subprocess.run(cmd, shell=True).returncode
        if i != 0: print(f'Error occured at running {cmd}')

def run(cmd: str):
    print(f'Running {cmd} at {Path.cwd()}')
    i = subprocess.run(cmd, shell=True).returncode
    if i != 0: print(f'Error occured at running {cmd}')

git_clone("https://github.com/comfyanonymous/ComfyUI", ".", True)
run('pip install -U onnxruntime-gpu')

get_file("https://huggingface.co/SG161222/RealVisXL_V4.0_Lightning/resolve/main/RealVisXL_V4.0_Lightning.safetensors", "./models/checkpoints/")

run(f"python main.py --listen 0.0.0.0 --port 7860 --cpu --output-directory {DATA_DIR}")

# https://github.com/comfyanonymous/ComfyUI