import os
import time
import sys
import fileinput
from IPython.display import clear_output
from subprocess import getoutput
from IPython.utils import capture
from pyngrok import ngrok, conf
import base64

blasphemy=base64.b64decode(("d2VidWk=")).decode('ascii')

Previous_Session="" #@param{type: 'string'}

#@markdown - Leave empty if you want to use the current trained model.

Use_Custom_Path = False #@param {type:"boolean"}

try:
  INSTANCE_NAME
  INSTANCET=INSTANCE_NAME
except:
  pass
#@markdown - if checked, an input box will ask the full path to a desired model.

if Previous_Session!="":
  INSTANCET=Previous_Session
  INSTANCET=INSTANCET.replace(" ","_")

if Use_Custom_Path:
  try:
    INSTANCET
    del INSTANCET
  except:
    pass

try:
  INSTANCET
  if Previous_Session!="":
    path_to_trained_model='/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/'+Previous_Session+"/"+Previous_Session+'.ckpt'
  else:
    path_to_trained_model=SESSION_DIR+"/"+INSTANCET+'.ckpt'
except:
  print('It seems that you did not perform training during this session or you chose to use a custom path,\nprovide the full path to the model (including the name of the model):\n')
  path_to_trained_model=input()

while not os.path.exists(path_to_trained_model):
   print("The model doesn't exist on you Gdrive, use the file explorer to get the path : ")
   path_to_trained_model=input()

fgitclone = "git clone --depth 1"

with capture.capture_output() as cap:
    if not os.path.exists('/content/gdrive/MyDrive'):
      !mkdir -p /content/gdrive/MyDrive

if not os.path.exists('/content/gdrive/MyDrive/sd/stablediffusion'):
    !wget -q -O /content/sd_rep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_rep.tar.zst
    !tar -C  /content/gdrive/MyDrive --zstd -xf /content/sd_rep.tar.zst
    !rm /content/sd_rep.tar.zst
    clear_output()

with capture.capture_output() as cap:
  %cd /content/gdrive/MyDrive/sd
  !git clone -q --branch master https://github.com/AUTOMATIC1111/stable-diffusion-$blasphemy
  %cd stable-diffusion-$blasphemy
  !mkdir cache
  !sed -i 's@~/.cache@/content/gdrive/MyDrive/sd/stable-diffusion-{blasphemy}/cache@' /usr/local/lib/python3.10/dist-packages/transformers/utils/hub.py

  clear_output()
  !git reset --hard
  time.sleep(1)
  !rm webui.sh
  !git pull
  !git fetch --unshallow
  !git checkout a9eab236d7e8afa4d6205127904a385b2c43bb24

with capture.capture_output() as cap:
  if not os.path.exists('/tools/node/bin/lt'):
    !npm install -g localtunnel

Ngrok_token = "" #@param {type:"string"}

#@markdown - Input your ngrok token if you want to use ngrok server.

Use_localtunnel = False #@param {type:"boolean"}

User = "" #@param {type:"string"}
Password= "" #@param {type:"string"}
#@markdown - Add credentials to your Gradio interface (optional).

auth=f"--gradio-auth {User}:{Password}"
if User =="" or Password=="":
  auth=""

with capture.capture_output() as cap:
  %cd modules
  !wget -q -O paths.py https://github.com/TheLastBen/fast-stable-diffusion/raw/5632d2ef7fffd940976538d270854ec4faf26855/AUTOMATIC1111_files/paths.py
  !wget -q -O extras.py https://github.com/AUTOMATIC1111/stable-diffusion-$blasphemy/raw/a9eab236d7e8afa4d6205127904a385b2c43bb24/modules/extras.py
  !wget -q -O sd_models.py https://github.com/AUTOMATIC1111/stable-diffusion-$blasphemy/raw/a9eab236d7e8afa4d6205127904a385b2c43bb24/modules/sd_models.py
  !wget -q -O /usr/local/lib/python3.10/dist-packages/gradio/blocks.py https://github.com/TheLastBen/fast-stable-diffusion/raw/7ff88eaa1fb4997bacd9845bd487f9a14335d625/AUTOMATIC1111_files/blocks.py
  %cd /content/gdrive/MyDrive/sd/stable-diffusion-$blasphemy/

  !sed -i "s@os.path.splitext(checkpoint_file)@os.path.splitext(checkpoint_file); map_location='cuda'@" /content/gdrive/MyDrive/sd/stable-diffusion-$blasphemy/modules/sd_models.py
  !sed -i 's@ui.create_ui().*@ui.create_ui();shared.demo.queue(concurrency_count=999999,status_update_rate=0.1)@' /content/gdrive/MyDrive/sd/stable-diffusion-$blasphemy/webui.py
  !sed -i "s@map_location='cpu'@map_location='cuda'@" /content/gdrive/MyDrive/sd/stable-diffusion-$blasphemy/modules/extras.py
  !sed -i 's@print(\"No module.*@@' /content/gdrive/MyDrive/sd/stablediffusion/ldm/modules/diffusionmodules/model.py
  !sed -i 's@\"quicksettings\": OptionInfo(.*@"quicksettings": OptionInfo("sd_model_checkpoint,  sd_vae, CLIP_stop_at_last_layers, inpainting_mask_weight, initial_noise_multiplier", "Quicksettings list"),@' /content/gdrive/MyDrive/sd/stable-diffusion-$blasphemy/modules/shared.py

share=''
if Ngrok_token!="":
  ngrok.kill()
  srv=ngrok.connect(7860, pyngrok_config=conf.PyngrokConfig(auth_token=Ngrok_token) , bind_tls=True).public_url

  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 = "{srv[8:]}"\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)

elif Use_localtunnel:
  with capture.capture_output() as cap:
    share=''
    %cd /content
    !nohup lt --port 7860 > srv.txt 2>&1 &
    time.sleep(2)
    !grep -o 'https[^ ]*' /content/srv.txt >srvr.txt
    time.sleep(2)
    srv= getoutput('cat /content/srvr.txt')

    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 = "{srv[8:]}"\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)

    !rm /content/srv.txt /content/srvr.txt
    %cd /content/gdrive/MyDrive/sd/stable-diffusion-$blasphemy

else:
  share='--share'

configf="--api --disable-safe-unpickle --enable-insecure-extension-access --no-half-vae --opt-sdp-attention --no-download-sd-model --disable-console-progressbars"

clear_output()

if os.path.isfile(path_to_trained_model):
  !python /content/gdrive/MyDrive/sd/stable-diffusion-$blasphemy/webui.py $share --ckpt "$path_to_trained_model" $auth $configf
else:
  !python /content/gdrive/MyDrive/sd/stable-diffusion-$blasphemy/webui.py $share --ckpt-dir "$path_to_trained_model" $auth $configf