import os
from IPython.display import clear_output
from subprocess import call, getoutput, Popen, run
import time
import ipywidgets as widgets
import requests
import sys
import fileinput
from torch.hub import download_url_to_file
from urllib.parse import urlparse, parse_qs, unquote
import re
import six
import gdown

from urllib.request import urlopen, Request
import tempfile
from tqdm import tqdm 
from bs4 import BeautifulSoup
import zipfile




def Deps(force_reinstall):
    # Определяем путь к файлу-маркеру
    #marker_file = '/workspace/setup_complete'
    # if not force_reinstall and os.path.exists(marker_file):
    #     ntbks()
    #     print('Modules and notebooks updated, dependencies already installed')
    #     os.environ['TORCH_HOME'] = '/workspace/cache/torch'
    #     os.environ['PYTHONWARNINGS'] = 'ignore'
    #     call("pip install controlnet_aux --root-user-action=ignore --disable-pip-version-check", shell=True, stdout=open('/dev/null', 'w'))
    #     call("pip install diskcache --root-user-action=ignore --disable-pip-version-check", shell=True, stdout=open('/dev/null', 'w'))
    #     call("pip install pillow-avif-plugin --root-user-action=ignore --disable-pip-version-check", shell=True, stdout=open('/dev/null', 'w'))
    #     call('pip install --root-user-action=ignore --disable-pip-version-check -qq gradio==3.41.2', shell=True, stdout=open('/dev/null', 'w'))
    #     call("pip install pyre-extensions==0.0.29 --root-user-action=ignore --disable-pip-version-check", shell=True, stdout=open('/dev/null', 'w'))
    #     clear_output()
    #     done()
    # else:
    call('pip install --root-user-action=ignore --disable-pip-version-check --no-deps -qq gdown PyWavelets numpy==1.23.5 accelerate==0.12.0 --force-reinstall', shell=True, stdout=open('/dev/null', 'w'))
    if os.path.exists('deps'):
        call("rm -r deps", shell=True)
    if os.path.exists('diffusers'):
        call("rm -r diffusers", shell=True)
    if not os.path.exists('cache'):
        call('mkdir cache', shell=True)   
    if not os.path.exists('deps'):
        call('mkdir deps', shell=True)
        os.chdir('deps')
        dwn("https://huggingface.co/spaces/stlaurentjr/RNPD/resolve/main/deps/rnpddeps-t2.tar.zst", "/workspace/deps/rnpddeps-t2.tar.zst", "Installing dependencies")
    else:
        os.chdir('deps')
        
    call('tar -C / --zstd -xf rnpddeps-t2.tar.zst', shell=True, stdout=open('/dev/null', 'w'))
    call("sed -i 's@~/.cache@/workspace/cache@' /usr/local/lib/python3.10/dist-packages/transformers/utils/hub.py", shell=True)
    os.chdir('/workspace')
    call("pip install pyre-extensions==0.0.29 --root-user-action=ignore --disable-pip-version-check", shell=True, stdout=open('/dev/null', 'w'))
    call("pip install controlnet_aux --root-user-action=ignore --disable-pip-version-check", shell=True, stdout=open('/dev/null', 'w'))
    call("pip install diskcache --root-user-action=ignore --disable-pip-version-check", shell=True, stdout=open('/dev/null', 'w'))
    call("pip install pillow-avif-plugin --root-user-action=ignore --disable-pip-version-check", shell=True, stdout=open('/dev/null', 'w'))
    call("git clone --depth 1 -q --branch main https://github.com/TheLastBen/diffusers", shell=True, stdout=open('/dev/null', 'w'))
    call('pip install --root-user-action=ignore --disable-pip-version-check -qq gradio==3.41.2', shell=True, stdout=open('/dev/null', 'w'))
    #call("rm -r deps", shell=True)
    os.chdir('/workspace')
    os.environ['TORCH_HOME'] = '/workspace/cache/torch'
    os.environ['PYTHONWARNINGS'] = 'ignore'
    call("sed -i 's@text = _formatwarnmsg(msg)@text =\"\"@g' /usr/lib/python3.10/warnings.py", shell=True)
    # Создаём файл-маркер после успешного выполнения установок
    #with open(marker_file, 'w') as f:
    #    f.write('Setup complete')
        
    clear_output()
    done()



def dwn(url, dst, msg):
    file_size = None
    req = Request(url, headers={"User-Agent": "torch.hub"})
    u = urlopen(req)
    meta = u.info()
    if hasattr(meta, 'getheaders'):
        content_length = meta.getheaders("Content-Length")
    else:
        content_length = meta.get_all("Content-Length")
    if content_length is not None and len(content_length) > 0:
        file_size = int(content_length[0])

    with tqdm(total=file_size, disable=False, mininterval=0.5,
              bar_format=msg+' |{bar:20}| {percentage:3.0f}%') as pbar:
        with open(dst, "wb") as f:
            while True:
                buffer = u.read(8192)
                if len(buffer) == 0:
                    break
                f.write(buffer)
                pbar.update(len(buffer))
            f.close()


def ntbks():

    os.chdir('/workspace')
    if not os.path.exists('Latest_Notebooks'):
        call('mkdir Latest_Notebooks', shell=True)
    else:
        call('rm -r Latest_Notebooks', shell=True)
        call('mkdir Latest_Notebooks', shell=True)
    os.chdir('/workspace/Latest_Notebooks')
    call('wget -q -i https://huggingface.co/datasets/TheLastBen/RNPD/raw/main/Notebooks.txt', shell=True)
    call('rm Notebooks.txt', shell=True)
    os.chdir('/workspace')


def repo(Huggingface_token_optional):

    from slugify import slugify
    from huggingface_hub import HfApi, CommitOperationAdd, create_repo
    
    os.chdir('/workspace')
    if Huggingface_token_optional!="":
       username = HfApi().whoami(Huggingface_token_optional)["name"]
       backup=f"https://huggingface.co/datasets/{username}/fast-stable-diffusion/resolve/main/sd_backup_rnpd.tar.zst"
       headers = {"Authorization": f"Bearer {Huggingface_token_optional}"}
       response = requests.head(backup, headers=headers)
       if response.status_code == 302:
          print('Restoring the SD folder...')
          open('/workspace/sd_backup_rnpd.tar.zst', 'wb').write(requests.get(backup, headers=headers).content)
          call('tar --zstd -xf sd_backup_rnpd.tar.zst', shell=True)
          call('rm sd_backup_rnpd.tar.zst', shell=True)
       else:
          print('Backup not found, using a fresh/existing repo...')
          time.sleep(2)
          if not os.path.exists('/workspace/sd/stablediffusiond'): #reset later
             call('wget -q -O sd_mrep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_mrep.tar.zst', shell=True)
             call('tar --zstd -xf sd_mrep.tar.zst', shell=True)
             call('rm sd_mrep.tar.zst', shell=True)        
          os.chdir('/workspace/sd')
          if not os.path.exists('stable-diffusion-webui'):
              call('git clone -q --depth 1 --branch master https://github.com/AUTOMATIC1111/stable-diffusion-webui', shell=True)
        
    else:
        print('Installing/Updating the repo...')
        os.chdir('/workspace')
        if not os.path.exists('/workspace/sd/stablediffusiond'): #reset later
           call('wget -q -O sd_mrep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_mrep.tar.zst', shell=True)
           call('tar --zstd -xf sd_mrep.tar.zst', shell=True)
           call('rm sd_mrep.tar.zst', shell=True)        

        os.chdir('/workspace/sd')
        if not os.path.exists('stable-diffusion-webui'):
            call('git clone -q --depth 1 --branch master https://github.com/AUTOMATIC1111/stable-diffusion-webui', shell=True)


    os.chdir('/workspace/sd/stable-diffusion-webui/')
    call('git reset --hard', shell=True)
    print('')
    call('git pull', shell=True)
    os.makedirs('/workspace/sd/stable-diffusion-webui/repositories', exist_ok=True)
    call('git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui-assets /workspace/sd/stable-diffusion-webui/repositories/stable-diffusion-webui-assets', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
    os.chdir('/workspace')
    clear_output()
    done()



def mdl(Original_Model_Version, Path_to_MODEL, MODEL_LINK):

    import gdown
    
    src=getsrc(MODEL_LINK)  

    if not os.path.exists("/workspace/sd/stable-diffusion-webui/models/VAE"):
        call("mkdir -p /workspace/sd/stable-diffusion-webui/models/VAE", shell=True)

    # Получаем список файлов в директории /workspace/auto-VAE
    files = os.listdir("/workspace/auto-VAE")

    for file in files:
        source_path = os.path.join("/workspace/auto-VAE", file)
        target_path = os.path.join(
            "/workspace/sd/stable-diffusion-webui/models/VAE", file
        )
        # Проверяем, существует ли уже символическая ссылка
        if not os.path.exists(target_path):
            call(f"ln -s {source_path} {target_path}", shell=True)

    if not os.path.exists('/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv1-5.ckpt'):
        call('ln -s /workspace/auto-models/* /workspace/sd/stable-diffusion-webui/models/Stable-diffusion', shell=True)
    
    if not os.path.exists("/workspace/sd/stable-diffusion-webui/models/Lora"):
        call("mkdir -p /workspace/sd/stable-diffusion-webui/models/Lora", shell=True)
        call(
            "ln -s /workspace/auto-lora/* /workspace/sd/stable-diffusion-webui/models/Lora",
            shell=True,
        )
        
    if Path_to_MODEL !='':
      if os.path.exists(str(Path_to_MODEL)):
        print('Using the custom model')
        model=Path_to_MODEL
      else:
          print('Wrong path, check that the path to the model is correct')

    elif MODEL_LINK !="":

      if src=='civitai':
         modelname=get_name(MODEL_LINK, False)
         model=f'/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}'
         if not os.path.exists(model):
            dwn(MODEL_LINK, model, 'Downloading the custom model')
            clear_output()
         else:
            print('Model already exists')
      elif src=='gdrive':
         modelname=get_name(MODEL_LINK, True)
         model=f'/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}'
         if not os.path.exists(model):
            gdown.download(url=MODEL_LINK, output=model, quiet=False, fuzzy=True)
            clear_output()
         else:
            print('Model already exists')
      else:
         modelname=os.path.basename(MODEL_LINK)
         model=f'/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}'
         if not os.path.exists(model):
            gdown.download(url=MODEL_LINK, output=model, quiet=False, fuzzy=True)
            clear_output()
         else:
            print('Model already exists')

      if os.path.exists(model) and os.path.getsize(model) > 1810671599:
        print('Model downloaded, using the custom model.')
      else:
        call('rm '+model, shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
        print('Wrong link, check that the link is valid')

    else:
        if Original_Model_Version == "v1.5":
           model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv1-5.ckpt"
           print('Using the original V1.5 model')
        elif Original_Model_Version == "v2-512":
            model='/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv2-512.ckpt'
            if not os.path.exists('/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv2-512.ckpt'):
               print('Downloading the V2-512 model...')
               call('gdown -O '+model+' https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-nonema-pruned.ckpt', shell=True)
               clear_output()
            print('Using the original V2-512 model')
        elif Original_Model_Version == "v2-768":
           model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv2-768.ckpt"
           print('Using the original V2-768 model')
        elif Original_Model_Version == "SDXL":
            model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/sd_xl_base_1.0.safetensors"
            print('Using the original SDXL model')

        else:
            model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion"
            print('Wrong model version, try again')
    try:
        model
    except:
        model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion"

    return model
    
def modeldwn(model_LINK):

    if model_LINK == "":
        print("[1;33mNothing to do")
    else:
        os.makedirs(
            "/workspace/sd/stable-diffusion-webui/models/Stable-diffusion",
            exist_ok=True,
        )

        src = getsrc(model_LINK)

        if src == "civitai":
            modelname = get_name(model_LINK, False)
            loramodel = f"/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}"
            if not os.path.exists(loramodel):
                dwn(model_LINK, loramodel, "Downloading the LoRA model")
                clear_output()
            else:
                print("[1;33mModel already exists")
        elif src == "gdrive":
            modelname = get_true_name(model_LINK)
            loramodel = f"/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}"
            if not os.path.exists(loramodel):
                gdown.download(
                    url=model_LINK.replace("/file/d/", "/uc?id=").replace("/view", ""),
                    output=loramodel,
                    quiet=False,
                )
                clear_output()
            else:
                print("[1;33mModel already exists")
        else:
            modelname = os.path.basename(model_LINK)
            loramodel = f"/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}"
            if not os.path.exists(loramodel):
                gdown.download(
                    url=model_LINK, output=loramodel, quiet=False, fuzzy=True
                )
                clear_output()
            else:
                print("[1;33mModel already exists")

        if os.path.exists(loramodel):
            print("[1;32mCheckpoints downloaded")
        else:
            print("[1;31mWrong link, check that the link is valid")

def loradwn(LoRA_LINK):

    if LoRA_LINK == "":
        print("[1;33mNothing to do")
    else:
        os.makedirs("/workspace/sd/stable-diffusion-webui/models/Lora", exist_ok=True)

        src = getsrc(LoRA_LINK)

        if src == "civitai":
            modelname = get_name(LoRA_LINK, False)
            loramodel = f"/workspace/sd/stable-diffusion-webui/models/Lora/{modelname}"
            if not os.path.exists(loramodel):
                dwn(LoRA_LINK, loramodel, "Downloading the LoRA model")
                clear_output()
            else:
                print("[1;33mModel already exists")
        elif src == "gdrive":
            modelname = get_true_name(LoRA_LINK)
            loramodel = f"/workspace/sd/stable-diffusion-webui/models/Lora/{modelname}"
            if not os.path.exists(loramodel):
                gdown.download(
                    url=LoRA_LINK.replace("/file/d/", "/uc?id=").replace("/view", ""),
                    output=loramodel,
                    quiet=False,
                )
                clear_output()
            else:
                print("[1;33mModel already exists")
        else:
            modelname = os.path.basename(LoRA_LINK)
            loramodel = f"/workspace/sd/stable-diffusion-webui/models/Lora/{modelname}"
            if not os.path.exists(loramodel):
                gdown.download(url=LoRA_LINK, output=loramodel, quiet=False, fuzzy=True)
                clear_output()
            else:
                print("[1;33mModel already exists")

        if os.path.exists(loramodel):
            print("[1;32mLoRA downloaded")
        else:
            print("[1;31mWrong link, check that the link is valid")

def CNet(ControlNet_Model, ControlNet_XL_Model):
    def download(url, model_dir):
        try:
            filename = os.path.basename(urlparse(url).path)
            pth = os.path.abspath(os.path.join(model_dir, filename))
            if not os.path.exists(pth):
                print("Downloading: " + os.path.basename(url))
                download_url_to_file(url, pth, hash_prefix=None, progress=True)
            else:
                print(f"[1;32mThe model {filename} already exists[0m")
        except Exception as e:
            print(f"Ошибка при скачивании {url}: {e}")

    wrngv1 = False
    os.chdir("/workspace/sd/stable-diffusion-webui/extensions")
    if not os.path.exists("sd-webui-controlnet"):
        call(
            "git clone https://github.com/Mikubill/sd-webui-controlnet.git", shell=True
        )
        os.chdir("/workspace")
    else:
        os.chdir("sd-webui-controlnet")
        call(
            "git reset --hard",
            shell=True,
            stdout=open("/dev/null", "w"),
            stderr=open("/dev/null", "w"),
        )
        call(
            "git pull",
            shell=True,
            stdout=open("/dev/null", "w"),
            stderr=open("/dev/null", "w"),
        )
        os.chdir("/workspace")

    mdldir = (
        "/workspace/sd/stable-diffusion-webui/extensions/sd-webui-controlnet/models"
    )
    call(
        "ln -s /workspace/auto-controlnet-models/* /workspace/sd/stable-diffusion-webui/extensions/sd-webui-controlnet/models",
        shell=True,
    )

    for filename in os.listdir(mdldir):
        if "_sd14v1" in filename:
            renamed = re.sub("_sd14v1", "-fp16", filename)
            os.rename(os.path.join(mdldir, filename), os.path.join(mdldir, renamed))

    call(
        "wget -q -O CN_models.txt https://huggingface.co/spaces/stlaurentjr/RNPD/raw/main/config/CN_models.txt",
        shell=True,
    )
    call(
        "wget -q -O CN_models_XL.txt https://github.com/TheLastBen/fast-stable-diffusion/raw/main/AUTOMATIC1111_files/CN_models_XL.txt",
        shell=True,
    )

    with open("CN_models.txt", "r") as f:
        mdllnk = f.read().splitlines()
    with open("CN_models_XL.txt", "r") as d:
        mdllnk_XL = d.read().splitlines()
    call("rm CN_models.txt CN_models_XL.txt", shell=True)

    os.chdir("/workspace")

    if ControlNet_Model == "All" or ControlNet_Model == "all":
        for lnk in mdllnk:
            download(lnk, mdldir)
        clear_output()

    elif ControlNet_Model == "17":
        mdllnk = list(filter(lambda x: "t2i" in x, mdllnk))
        for lnk in mdllnk:
            download(lnk, mdldir)
        clear_output()

    elif (
        ControlNet_Model.isdigit()
        and int(ControlNet_Model) - 1 < 16
        and int(ControlNet_Model) > 0
    ):
        download(mdllnk[int(ControlNet_Model) - 1], mdldir)
        clear_output()

    elif ControlNet_Model == "none":
        pass
        clear_output()

    else:
        print("[1;31mWrong ControlNet V1 choice, try again")
        wrngv1 = True

    if ControlNet_XL_Model == "All" or ControlNet_XL_Model == "all":
        for lnk_XL in mdllnk_XL:
            download(lnk_XL, mdldir)
        if not wrngv1:
            clear_output()
        done()

    elif ControlNet_XL_Model.isdigit() and int(ControlNet_XL_Model) - 1 < 5:
        download(mdllnk_XL[int(ControlNet_XL_Model) - 1], mdldir)
        if not wrngv1:
            clear_output()
        done()

    elif ControlNet_XL_Model == "none":
        pass
        if not wrngv1:
            clear_output()
        done()

    else:
        print("[1;31mWrong ControlNet XL choice, try again")

def sd(User, Password, model):

    import gradio
    
    gradio.close_all()
    
    auth=f"--gradio-auth {User}:{Password}"
    if User =="" or Password=="":
      auth=""

    call('wget -q -O /usr/local/lib/python3.10/dist-packages/gradio/blocks.py https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/AUTOMATIC1111_files/blocks.py', shell=True)
   
    os.chdir('/workspace/sd/stable-diffusion-webui/modules')

    call("sed -i 's@possible_sd_paths =.*@possible_sd_paths = [\"/workspace/sd/stablediffusion\"]@' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
    call("sed -i 's@\.\.\/@src/@g' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
    call("sed -i 's@src\/generative-models@generative-models@g' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
    
    call("sed -i 's@\[\"sd_model_checkpoint\"\]@\[\"sd_model_checkpoint\", \"sd_vae\", \"CLIP_stop_at_last_layers\", \"inpainting_mask_weight\", \"initial_noise_multiplier\"\]@g' /workspace/sd/stable-diffusion-webui/modules/shared.py", shell=True)
    
    call("sed -i 's@print(\"No module.*@@' /workspace/sd/stablediffusion/ldm/modules/diffusionmodules/model.py", shell=True)
    os.chdir('/workspace/sd/stable-diffusion-webui')
    clear_output()

    podid=os.environ.get('RUNPOD_POD_ID')
    localurl=f"{podid}-3001.proxy.runpod.net"

    for line in fileinput.input('/usr/local/lib/python3.10/dist-packages/gradio/blocks.py', inplace=True):
      if line.strip().startswith('self.server_name ='):
          line = f'            self.server_name = "{localurl}"\n'
      if line.strip().startswith('self.protocol = "https"'):
          line = '            self.protocol = "https"\n'
      if line.strip().startswith('if self.local_url.startswith("https") or self.is_colab'):
          line = ''
      if line.strip().startswith('else "http"'):
          line = ''
      sys.stdout.write(line)
      
    if model=="":
        mdlpth=""
    else:
        if os.path.isfile(model):
            mdlpth="--ckpt "+model
        else:
            mdlpth="--ckpt-dir "+model
    vae_path = "--vae-path /workspace/sd/stable-diffusion-webui/models/VAE/vae-ft-mse-840000-ema-pruned.safetensors"
    configf="--disable-console-progressbars --no-half-vae --disable-safe-unpickle --api --no-download-sd-model --opt-sdp-attention --enable-insecure-extension-access  --skip-version-check --listen --port 3000 "+auth+" "+mdlpth+" "+vae_path

    return configf



def save(Huggingface_Write_token):
    
    from slugify import slugify
    from huggingface_hub import HfApi, CommitOperationAdd, create_repo
    
    if Huggingface_Write_token=="":
        print('A huggingface write token is required')
        
    else:
        os.chdir('/workspace')
        
        if os.path.exists('sd'):
        
            call('tar --exclude="stable-diffusion-webui/models/*/*" --exclude="sd-webui-controlnet/models/*" --zstd -cf sd_backup_rnpd.tar.zst sd', shell=True)
            api = HfApi()
            username = api.whoami(token=Huggingface_Write_token)["name"]

            repo_id = f"{username}/{slugify('fast-stable-diffusion')}"

            print("Backing up...")

            operations = [CommitOperationAdd(path_in_repo="sd_backup_rnpd.tar.zst", path_or_fileobj="/workspace/sd_backup_rnpd.tar.zst")]

            create_repo(repo_id,private=True, token=Huggingface_Write_token, exist_ok=True, repo_type="dataset")

            api.create_commit(
              repo_id=repo_id,
              repo_type="dataset",
              operations=operations,
              commit_message="SD folder Backup",
              token=Huggingface_Write_token
            )

            call('rm sd_backup_rnpd.tar.zst', shell=True)
            clear_output()

            done()
     
        else:
            print('Nothing to backup')
        
def download_and_install_config():
    target_directory = "/workspace/sd/stable-diffusion-webui/"
    config_url = (
        "https://huggingface.co/spaces/stlaurentjr/RNPD/raw/main/config/config.json"
    )
    config_file_path = os.path.join(target_directory, "config.json")

    # Check if the config file already exists
    if not os.path.exists(config_file_path):
        # Change the directory
        os.chdir(target_directory)
        # Download the file using curl command
        call(f"curl -o config.json {config_url}", shell=True)
        print(f"Config file downloaded successfully.")
    else:
        print("Config file already exists, download skipped.")


def InstallDeforum():
    # Переход в директорию расширений
    os.chdir("/workspace/sd/stable-diffusion-webui/extensions")

    # Проверка наличия директории с расширением
    if not os.path.exists("deforum-for-automatic1111-webui"):
        # Клонирование репозитория, если он не существует
        call(
            "git clone https://github.com/deforum-art/sd-webui-deforum.git", shell=True
        )
        # Возврат в корневую директорию
        os.chdir("/workspace")
    else:
        # Если директория существует, переходим в нее
        # os.chdir('deforum-for-automatic1111-webui')
        # Сброс изменений и обновление репозитория
        # call('git reset --hard', shell=True, stdout=open(os.devnull, 'w'), stderr=open(os.devnull, 'w'))
        # call('git pull', shell=True, stdout=open(os.devnull, 'w'), stderr=open(os.devnull, 'w'))
        # Возврат в корневую директорию
        os.chdir("/workspace")
        
def InstallReActor():
    # Переход в директорию расширений
    os.chdir("/workspace/sd/stable-diffusion-webui/extensions")
    # Проверка наличия директории с расширением
    if not os.path.exists("sd-webui-reactor"):
        # Установка библиотек
        from subprocess import call, DEVNULL
        call(
            'pip install "onnxruntime-gpu>=1.16.1"', 
             shell=True
        )
        # Клонирование репозитория, если он не существует
        call(
            "git clone https://github.com/Gourieff/sd-webui-reactor.git",
            shell=True,
        )
        # Создание папки под модели
        call(
            "mkdir /workspace/sd/stable-diffusion-webui/models/insightface",
            shell=True
        )
        os.chdir("/workspace/sd/stable-diffusion-webui/models/insightface")
        # Скачивание модели для reActor
        call(
            "wget https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx", 
            shell=True, 
            stdout=DEVNULL, 
            stderr=DEVNULL
        )
        # Вовзрат в корневую директорию
        clear_output()
        os.chdir("/workspace")
    else:
        #Возврат в корневую директорию
        os.chdir("/workspace")

def InstallAnimateDiff():
    # Переход в директорию расширений
    os.chdir("/workspace/sd/stable-diffusion-webui/extensions")

    # Проверка наличия директории с расширением
    if not os.path.exists("sd-webui-animatediff"):
        from subprocess import call, DEVNULL
        # Клонирование репозитория, если он не существует
        call(
            "git clone https://github.com/continue-revolution/sd-webui-animatediff.git",
            shell=True,
        )
        # Переход в папку с моделями
        os.chdir(
            "/workspace/sd/stable-diffusion-webui/extensions/sd-webui-animatediff/model"
        )
        # Скачивание моделей для animatediff
        call(
            "wget https://huggingface.co/guoyww/animatediff/resolve/refs%2Fpr%2F3/mm_sd_v15_v2.ckpt",
            shell=True, 
            stdout=DEVNULL, 
            stderr=DEVNULL
        )
      #  call(
      #      "wget https://huggingface.co/manshoety/AD_Stabilized_Motion/resolve/main/mm-Stabilized_high.pth",
      #      shell=True, 
      #      stdout=DEVNULL, 
      #      stderr=DEVNULL
      #  )
        # Вовзрат в корневую директорию
        clear_output()
        os.chdir("/workspace")
    else:
        # Возврат в корневую директорию
        os.chdir("/workspace")

def InstallWildCards():
    # Путь к целевой директории
    target_dir = "/workspace/sd/stable-diffusion-webui/extensions/stable-diffusion-webui-wildcards/wildcards/"

    # Перемещаемся в директорию расширений
    os.chdir("/workspace/sd/stable-diffusion-webui/extensions")

    if not os.path.exists("stable-diffusion-webui-wildcards"):
        # Клонирование репозитория, если он ещё не существует
        call(
            "git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui-wildcards.git",
            shell=True,
        )

        # Перемещаемся в клонированный репозиторий
        os.chdir("stable-diffusion-webui-wildcards")

        # Скачиваем файл по ID
        gdown.download(
            id="1sY9Yv29cCYZuxBvszkmLVgw--aRdwT1P", output="wildcards.zip", quiet=False
        )

        # Создаем целевую директорию, если она ещё не существует
        os.makedirs(target_dir, exist_ok=True)

        # Распаковываем архив
        with zipfile.ZipFile("wildcards.zip", "r") as zip_ref:
            zip_ref.extractall(target_dir)
    else:
        # Если репозиторий уже существует, обновляем его
        os.chdir("stable-diffusion-webui-wildcards")
        call(
            "git reset --hard",
            shell=True,
            stdout=open(os.devnull, "w"),
            stderr=open(os.devnull, "w"),
        )
        call(
            "git pull",
            shell=True,
            stdout=open(os.devnull, "w"),
            stderr=open(os.devnull, "w"),
        )

    # Возвращаемся в исходную рабочую директорию
    os.chdir("/workspace")
def getsrc(url):

    parsed_url = urlparse(url)

    if parsed_url.netloc == "civitai.com":
        src = "civitai"
    elif parsed_url.netloc == "drive.google.com":
        src = "gdrive"
    elif parsed_url.netloc == "huggingface.co":
        src = "huggingface"
    else:
        src = "others"
    return src


def get_true_name(url):
    response = requests.get(url)
    soup = BeautifulSoup(response.text, "html.parser")
    title_tag = soup.find("title")
    if title_tag:
        title_text = title_tag.text
        # Извлечение имени файла из тега title (предполагая, что имя файла указано в теге title)
        file_name = title_text.split(" - ")[0]
        return file_name
    else:
        raise RuntimeError("Could not find the title tag in the HTML")


def get_name(url, gdrive):

    from gdown.download import get_url_from_gdrive_confirmation

    if not gdrive:
        response = requests.get(url, allow_redirects=False)
        if "Location" in response.headers:
            redirected_url = response.headers["Location"]
            quer = parse_qs(urlparse(redirected_url).query)
            if "response-content-disposition" in quer:
                disp_val = quer["response-content-disposition"][0].split(";")
                for vals in disp_val:
                    if vals.strip().startswith("filename="):
                        filenm = unquote(vals.split("=", 1)[1].strip())
                        return filenm.replace('"', "")
    else:
        headers = {
            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36"
        }
        lnk = "https://drive.google.com/uc?id={id}&export=download".format(
            id=url[url.find("/d/") + 3 : url.find("/view")]
        )
        res = requests.session().get(lnk, headers=headers, stream=True, verify=True)
        res = requests.session().get(
            get_url_from_gdrive_confirmation(res.text),
            headers=headers,
            stream=True,
            verify=True,
        )
        content_disposition = six.moves.urllib_parse.unquote(
            res.headers["Content-Disposition"]
        )
        filenm = (
            re.search(r"filename\*=UTF-8''(.*)", content_disposition)
            .groups()[0]
            .replace(os.path.sep, "_")
        )
        return filenm


def done():
    done = widgets.Button(
        description='Done!',
        disabled=True,
        button_style='success',
        tooltip='',
        icon='check'
    )
    display(done)