auto_ncpt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# NCPT Paperspace One-click (v0.3)\n",
    "\n",
    "### Last updated: 25 Dec 2023\n",
    "\n",
    "Check for updates [here](https://github.com/NoCrypt/paperspace-ncpt)\n",
    "\n",
    "**!! ONLY RUN IN JUPYTERLAB !!**\n",
    "\n",
    "![](https://i.ibb.co/M9ZH6qf/image.png)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "To hide the code, press on this blue bar on the left of the cell\n",
    "\n",
    "![](https://i.ibb.co/gFqKVV6/chrome-23-12-25-195649.png)\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# fix gdown issue by downgrading it \n",
    "!pip uninstall gdown -y\n",
    "!pip install 'gdown==4.6.0'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Run to initialize and set parameters    . \n",
    "import site, sys, time, json, glob, zipfile, re, time, shutil, gdown, subprocess, urllib\n",
    "from urllib.parse import urlparse, parse_qs, unquote\n",
    "from typing import List, Tuple, Optional, Any\n",
    "from ipywidgets import Checkbox, Text, Button, HBox, ToggleButton, Label, Box, Layout, Accordion, VBox, CallbackDispatcher\n",
    "from IPython.display import display, HTML, clear_output\n",
    "from pathlib import Path\n",
    "from subprocess import check_output, check_call\n",
    "from tempfile import NamedTemporaryFile\n",
    "from os import getenv, environ\n",
    "\n",
    "# Mapping resources folders \n",
    "paths = {\n",
    "    \"models\": \"models/Stable-diffusion\",\n",
    "    \"loras\": \"models/Lora\",\n",
    "    \"vaes\": \"models/VAE\",\n",
    "    \"embeddings\": \"embeddings\",\n",
    "    \"extensions\":  \"extensions\",  \n",
    "}\n",
    "\n",
    "# options for aria2 to use\n",
    "aria2_opts = [ \"--console-log-level=error\", \"--summary-interval=30\", \"-j1\", \"-x16\", \"-s16\", \"-k1M\", \"-c\"]\n",
    "\n",
    "\n",
    "# --- various methods --- #\n",
    "\n",
    "\n",
    "def filename_wget(url :str, timeout : int=5) -> Optional[str]:\n",
    "\n",
    "    #batchlinks-downloader method\n",
    "    req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})\n",
    "    response = urllib.request.urlopen(req, timeout=timeout)\n",
    "    forward_url = response.geturl()\n",
    "\n",
    "    match = re.search(r'filename%3D%22(.+?)%22', forward_url)\n",
    "    if match:\n",
    "        filename = urllib.parse.unquote(match.group(1))\n",
    "        return filename\n",
    "\n",
    "    #wget method\n",
    "    dst=Path('/notebooks/.tmp')\n",
    "    if not dst.exists():\n",
    "        dst.mkdir()\n",
    "        \n",
    "    process = subprocess.Popen(['wget', '--content-disposition', '-P', dst, url], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n",
    "    \n",
    "    filename = None\n",
    "    start_time = time.time()\n",
    "\n",
    "    while True:\n",
    "        if time.time() - start_time > timeout:\n",
    "            display(\"Timeout reached.\")\n",
    "            break\n",
    "        line = process.stdout.readline()\n",
    "        if line == '' and process.poll() is not None:\n",
    "            break\n",
    "        if 'Saving to:' in line:\n",
    "            match = re.search(r'Saving to: ‘.+/(.+)’', line)\n",
    "            if match:\n",
    "                filename = match.group(1)\n",
    "                break\n",
    "    if process.poll() is None:\n",
    "        process.terminate()\n",
    "        try:\n",
    "            process.wait(timeout=0.5)\n",
    "        except subprocess.TimeoutExpired:\n",
    "            process.kill()\n",
    "    if filename:\n",
    "        file = dst.joinpath(filename)\n",
    "        if file.exists() :\n",
    "            file.unlink()\n",
    "    return filename\n",
    "    \n",
    "def delete_symlinks(directory_path:str)  -> None :\n",
    "    directory = Path(directory_path)\n",
    "    if not directory.exists() or not directory.is_dir():\n",
    "        return\n",
    "    for item in directory.iterdir():  \n",
    "        if item.is_symlink():\n",
    "            item.unlink() \n",
    " \n",
    "# Folder size method \n",
    "def get_directory_size(directory_pattern:str) -> str:\n",
    "    try:\n",
    "        if '*' in directory_pattern:\n",
    "            directories = glob.glob(directory_pattern)\n",
    "        else:\n",
    "            directory_pattern_path = Path(directory_pattern)\n",
    "            directories = [directory_pattern] if directory_pattern_path.is_dir() else []\n",
    "        if not directories:\n",
    "            return \"0B\"\n",
    "        total_size = 0\n",
    "        for directory in directories:\n",
    "            output = check_output(['du', '-s', directory], text=True)\n",
    "            size = int(output.split()[0])  # Size in blocks\n",
    "            total_size += size * 1024  # Convert to bytes (assuming block size is 1024 bytes)\n",
    "\n",
    "        # Convert bytes to a human-readable format\n",
    "        for unit in ['B', 'KB', 'MB', 'GB', 'TB', 'PB']:\n",
    "            if total_size < 1024:\n",
    "                return f\"{total_size:.2f}{unit}\"\n",
    "            total_size /= 1024\n",
    "\n",
    "        return f\"{total_size:.2f}PB\"\n",
    "    except Exception as e:\n",
    "        #print(\"An error occurred:\", e)\n",
    "        return \"0B\"\n",
    "\n",
    "# copy dep folder\n",
    "def copy_recent_files(src_dir:str, dst_dir:str, minutes:int=10) -> None:\n",
    "    # Calculate the cutoff time\n",
    "    cutoff_time = time.time() - minutes * 60 *1000\n",
    "    \n",
    "    # Ensure the destination directory exists\n",
    "    dst_dir_path = Path(dst_dir)\n",
    "    dst_dir_path.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    # Iterate through the immediate files in src_dir\n",
    "    for item in Path(src_dir).iterdir():\n",
    "        if item.stat().st_mtime >= cutoff_time:\n",
    "            if item.is_dir():\n",
    "                # Use shutil.copytree for directories\n",
    "                shutil.copytree(\n",
    "                    item,\n",
    "                    dst_dir_path / item.name,\n",
    "                    symlinks=True,\n",
    "                    ignore=shutil.ignore_patterns(\"*.pyc\"),\n",
    "                    dirs_exist_ok=True,\n",
    "                )\n",
    "            elif item.is_file() and not item.name.endswith('.pyc'):\n",
    "                # Use shutil.copy2 for files\n",
    "                shutil.copy2(item, dst_dir_path / item.name)\n",
    "                \n",
    "# get current dep folder path\n",
    "def get_python_package_path()  -> str:\n",
    "    # Try using site.getsitepackages() first\n",
    "    package_paths = site.getsitepackages()\n",
    "    # Add the user site-package directory\n",
    "    package_paths.append(site.getusersitepackages())\n",
    "\n",
    "    # Check if these paths exist\n",
    "    for path_str in package_paths:\n",
    "        path = Path(path_str)\n",
    "        if path.exists():\n",
    "            return str(path)\n",
    "\n",
    "    # Fallback: Construct the path using sys.executable\n",
    "    python_executable_dir = Path(sys.executable).parent\n",
    "    for dir_name in ['site-packages', 'dist-packages']:\n",
    "        fallback_path = python_executable_dir / dir_name\n",
    "        if fallback_path.exists():\n",
    "            return str(fallback_path)\n",
    "    raise FileNotFoundError(\"Python package directory not found\")\n",
    "\n",
    "\n",
    "def wget_name(url:str,dst:str) -> Optional[str]:\n",
    "    # Running wget command to get headers\n",
    "    output = !wget -S -P {dst} --content-disposition {url}\n",
    "    # Process the output to extract the filename\n",
    "    output_str = \"\\n\".join(output)\n",
    "    # Use regular expression to find the filename\n",
    "    match = re.search(r'Content-Disposition:.*filename=\"(.+)\"', output_str, re.IGNORECASE)\n",
    "\n",
    "    if match :\n",
    "        new_file=f\"{dst}/{match.group(1)}\"\n",
    "        #cleaning up the uwanted duplicates if there is any\n",
    "        !rm -rf {new_file}.*\n",
    "        return match.group(1)\n",
    "    return None\n",
    "\n",
    "\n",
    "class Params:\n",
    "    # default parameters\n",
    "    file_path=\"\"\n",
    "    config_file=\"config.json\"\n",
    "    controls =[]\n",
    "    accordions ={}\n",
    "    links={\n",
    "        \"models\": [\"https://huggingface.co/NoCrypt/expermental_models/resolve/main/shux7.safetensors\"],\n",
    "        \"vaes\":[\"https://huggingface.co/NoCrypt/resources/resolve/main/VAE/any.vae.safetensors\"],\n",
    "        \"loras\": [],\n",
    "        \"extensions\": [],\n",
    "        \"embeddings\": [],\n",
    "        \"ui\": {\"params\": {},\"enabled\": [\"https://huggingface.co/NoCrypt/expermental_models/resolve/main/shux7.safetensors\",\"https://huggingface.co/NoCrypt/resources/resolve/main/VAE/any.vae.safetensors\"],\"keep\": [],\"names\": {},\"tabs\": []}}\n",
    "\n",
    "    def __init__(self, control_list:List[str]=[]):\n",
    "        self.controls = control_list\n",
    "\n",
    "    def load(self, file:str=\"links.json\")  -> None:\n",
    "        self.file_path = Path(file)\n",
    "\n",
    "        if not self.file_path.exists():\n",
    "            self.save()\n",
    "        else:\n",
    "            with self.file_path.open('r') as file:\n",
    "                self.links = json.load(file)\n",
    "\n",
    "        for ctrl in self.controls:\n",
    "            if ctrl.tooltip in self.links[\"ui\"][\"params\"]:\n",
    "                ctrl.value = self.links[\"ui\"][\"params\"][ctrl.tooltip]\n",
    "\n",
    "    # update parameters from ui \n",
    "    def update(self)  -> None:\n",
    "        for ctrl in self.controls:\n",
    "            self.links[\"ui\"][\"params\"][ctrl.tooltip] = ctrl.value\n",
    "            \n",
    "    # save parameters to file           \n",
    "    def save(self) -> None:\n",
    "        self.update()\n",
    "        with self.file_path.open('w')  as file:\n",
    "            json.dump(self.links, file, indent=4)\n",
    "\n",
    "    # get the status of the ressource enabled/cached\n",
    "    def get_link_status(self, url:str) -> Tuple[bool, bool]:\n",
    "        enabled = url in self.links[\"ui\"][\"enabled\"]\n",
    "        keep = url in self.links[\"ui\"][\"keep\"]\n",
    "        return enabled, keep\n",
    "\n",
    "    # get the ressource name from parameters\n",
    "    def name_from_link(self , url:str)-> str:\n",
    "        if 'names' in self.links[\"ui\"] and url in self.links[\"ui\"]['names']:\n",
    "            return self.links[\"ui\"]['names'][url]\n",
    "        else:\n",
    "            return \"\"\n",
    "    def link_from_name(self, name: str) -> str:\n",
    "        if 'names' in self.links[\"ui\"]:\n",
    "            for url, n in self.links[\"ui\"]['names'].items():\n",
    "                if n == name:\n",
    "                    return url\n",
    "        return \"\"\n",
    "    \n",
    "    # get the ressource name from parameters\n",
    "    def set_link_name(self , url:str, name:str) -> None:\n",
    "        if 'names' not in self.links[\"ui\"] :\n",
    "            self.links[\"ui\"]['names']={}\n",
    "        self.links[\"ui\"]['names'][url]=name\n",
    "        \n",
    "    # update the ressource status\n",
    "    def update_link_status(self, url:str, to_enable:bool, to_keep:bool) -> None:\n",
    "        enabled, keep = self.get_link_status(url)\n",
    "        if to_enable :\n",
    "            if not enabled :\n",
    "                self.links[\"ui\"][\"enabled\"].append(url)\n",
    "        else :\n",
    "            if enabled :\n",
    "                self.links[\"ui\"][\"enabled\"].remove(url)\n",
    "                \n",
    "        if to_keep :\n",
    "            if not keep :\n",
    "                self.links[\"ui\"][\"keep\"].append(url)\n",
    "        else :\n",
    "            if keep :\n",
    "                self.links[\"ui\"][\"keep\"].remove(url)\n",
    "        self.save()\n",
    "\n",
    "    # adds a new ressource\n",
    "    def add_link(self, url:str, category:str)-> bool:\n",
    "        def is_valid_url(url):\n",
    "            parsed = urlparse(url)\n",
    "            return bool(parsed.scheme) and bool(parsed.netloc)\n",
    "            \n",
    "        url=url.strip()\n",
    "        if is_valid_url(url) and  url not in self.links[category] : \n",
    "            self.links[category].append(url)\n",
    "            self.links[\"ui\"][\"enabled\"].append(url)\n",
    "            \n",
    "            filename=self.links[\"ui\"][\"names\"].get(url,'')\n",
    "            if filename ==\"\" :\n",
    "                 filename=filename_wget(url)     \n",
    "            if filename :\n",
    "                self.links[\"ui\"][\"names\"][url]=filename if \"login\" not in filename else f\"[login error]{url}\"\n",
    "            self.get_accordion(category)\n",
    "            self.save()\n",
    "            return True\n",
    "        return False\n",
    "\n",
    "        \n",
    "        \n",
    "    # removes the specified ressource\n",
    "    def remove_link(self, category:str, url:str) -> None:\n",
    "        if url in self.links[category] :\n",
    "            self.links[category].remove(url)\n",
    "        if url in self.links[\"ui\"][\"enabled\"] :\n",
    "            self.links[\"ui\"][\"enabled\"].remove(url)\n",
    "        if url in self.links[\"ui\"][\"keep\"] :\n",
    "            self.links[\"ui\"][\"keep\"].remove(url)\n",
    "        if url in self.links[\"ui\"][\"names\"] :\n",
    "            del self.links[\"ui\"][\"names\"][url]\n",
    "        self.get_accordion(category) \n",
    "        self.save()\n",
    "\n",
    "\n",
    "    # create or update the ressource manager \n",
    "    def get_accordion(self, category:str, refresh:bool=False) -> Accordion:\n",
    "        is_new=False\n",
    "        if category not in self.accordions:\n",
    "            is_new=True\n",
    "            self.accordions[category] = Accordion(layout=Layout(width='52%'))\n",
    "            def on_accordion_change(change):\n",
    "                if \"tabs\" not in self.links[\"ui\"] :\n",
    "                    self.links[\"ui\"][\"tabs\"]=[]\n",
    "                \n",
    "                if change['name'] == 'selected_index' and change['new'] is not None:\n",
    "                    if category not in self.links[\"ui\"][\"tabs\"]:\n",
    "                        self.links[\"ui\"][\"tabs\"].append(category)\n",
    "                else :\n",
    "                    if category in self.links[\"ui\"][\"tabs\"]:\n",
    "                        self.links[\"ui\"][\"tabs\"].remove(category)\n",
    "            self.accordions[category].observe(on_accordion_change, names='selected_index')\n",
    "\n",
    "        if (len(self.accordions[category].children)-1) !=len(self.links[category]) or refresh or len(self.links[category])==0:\n",
    "            vbox_children = []\n",
    "            vbox_children.append(self.create_add_link(category))\n",
    "            for url in self.links[category]:\n",
    "                element_name=self.name_from_link(url)\n",
    "                if element_name == \"\" :\n",
    "                    element_name=url\n",
    "                hbox = HBox(list(self.create_action_buttons(url, category)) + [Label(value=element_name)])\n",
    "                vbox_children.append(hbox)\n",
    "  \n",
    "            self.accordions[category].children = [VBox(vbox_children)]\n",
    "    \n",
    "        self.update_accordion_title(category)\n",
    "        if is_new :\n",
    "            if \"tabs\" in self.links[\"ui\"] and self.links[\"ui\"][\"tabs\"]:\n",
    "                if category in self.links[\"ui\"][\"tabs\"] :\n",
    "                    self.accordions[category].selected_index = 0\n",
    "                else:\n",
    "                    self.accordions[category].selected_index = None\n",
    "\n",
    "        return self.accordions[category]\n",
    "    \n",
    "    # update the count of resources \n",
    "    def update_accordion_title(self, category : str) -> None:\n",
    "        # Display the update button at the end\n",
    "        keep_cpt=0\n",
    "        enabled_cpt=0\n",
    "        for url in self.links[category] :\n",
    "            if url in self.links[\"ui\"][\"keep\"] :\n",
    "                keep_cpt=keep_cpt+1 \n",
    "            if url in self.links[\"ui\"][\"enabled\"] :\n",
    "                enabled_cpt=enabled_cpt+1\n",
    "        self.accordions[category].set_title(0, str(len(self.links[category])) +\" \"+ category + \" [\"+str(enabled_cpt)+\"/\" +str(keep_cpt) +\"]\")\n",
    "        \n",
    "    #---- ui components ----# \n",
    "    # create action buttons for the link manager\n",
    "    def create_action_buttons(self, url :str , category : str) -> Tuple[ToggleButton, ToggleButton, Button]:\n",
    "        enabled, keep = self.get_link_status(url)\n",
    "        \n",
    "        toggle_enable =ToggleButton(\n",
    "            layout=Layout(width='90px'),\n",
    "            value=enabled, \n",
    "            description=\"Enable\",\n",
    "            tooltip=\"Enable this model\",\n",
    "            button_style=('success' if enabled else ''),\n",
    "        )\n",
    "    \n",
    "        toggle_keep =ToggleButton(\n",
    "            layout=Layout(width='90px'),\n",
    "            value=keep, \n",
    "            description=\"Cache\",\n",
    "            tooltip=\"Cache this model into notebooks\",\n",
    "            button_style=('success' if keep else ''),\n",
    "        )\n",
    "    \n",
    "        delete_button = Button(\n",
    "            tooltip=\"Delete this model\",\n",
    "            button_style='danger',     \n",
    "            layout=Layout(width='30px'),\n",
    "            icon='trash-o' \n",
    "        )\n",
    "        \n",
    "        def on_toggle(cc:CallbackDispatcher):\n",
    "            self.update_link_status(url,toggle_enable.value,toggle_keep.value)\n",
    "            toggle_keep.button_style = 'success' if toggle_keep.value else ''\n",
    "            toggle_enable.button_style = 'success' if toggle_enable.value else ''\n",
    "            self.update_accordion_title(category)\n",
    "            self.save()\n",
    "        \n",
    "        def on_delete(bb:Button):\n",
    "            self.update_link_status(url, False, False)\n",
    "            self.remove_link(category,url)\n",
    "            self.save()\n",
    "\n",
    "        toggle_enable.observe(on_toggle, 'value')\n",
    "        toggle_keep.observe(on_toggle, 'value')\n",
    "        delete_button.on_click(on_delete)\n",
    "        return toggle_enable , toggle_keep , delete_button\n",
    "\n",
    "    # create add ressource ui\n",
    "    def create_add_link(self, category:str)-> HBox:\n",
    "        link_control = Text(\n",
    "            value=\"\",\n",
    "            placeholder=category + \" link...\",\n",
    "            layout=Layout(width='80%'),\n",
    "        )\n",
    "        \n",
    "        add_button = Button(\n",
    "            description='Add',\n",
    "            layout=Layout(width='20%'),\n",
    "            icon='plus' \n",
    "        )\n",
    "    \n",
    "        def on_add(bb:Button) -> None:            \n",
    "            description='Adding ...'\n",
    "            add_button.button_style=\"warning\"\n",
    "            add_button.icon=\"refresh\"\n",
    "            if self.add_link(link_control.value,category ) :\n",
    "                add_button.button_style=\"\"\n",
    "                add_button.icon=\"plus\"\n",
    "                description='Add'\n",
    "            else :\n",
    "                add_button.button_style=\"danger\"\n",
    "                add_button.icon=\"plus\"\n",
    "                description='Add'\n",
    "                \n",
    "        add_button.on_click(on_add)\n",
    "        return HBox([link_control, add_button])\n",
    "        \n",
    "    # creating output zip/clean tools \n",
    "    def create_zip_tool(self) -> HBox:\n",
    "        #view\n",
    "        zip_button = Button(\n",
    "            description='Zip',\n",
    "            tooltip='Zip folder',\n",
    "            icon='cube' \n",
    "        )\n",
    "        clean_button = Button(\n",
    "            description='Clean',\n",
    "            tooltip='Cleaning folder',\n",
    "            icon='trash-o'\n",
    "        )\n",
    "        zip_clean_button = Button(\n",
    "            description='Zip & Clean',\n",
    "            button_style='info', \n",
    "            tooltip='Zip folder and clean it up',\n",
    "            icon='codepen'\n",
    "        )\n",
    "            \n",
    "        def on_zip(bb:Button):\n",
    "            display(\"Cleaning..a.\")\n",
    "            bb.description = bb.description.split('>')[0]\n",
    "            is_cleaning = \"clean\" in bb.description.lower()\n",
    "            is_zipping = \"zip\" in bb.description.lower()\n",
    "            cleanup_dir = Path(output_control.value)\n",
    "            if not cleanup_dir.exists():\n",
    "                display(\"The folder doesn't exist.\")\n",
    "                bb.button_style = 'danger'  \n",
    "                bb.icon=\"times\"\n",
    "                \n",
    "            if any(cleanup_dir.iterdir()) :\n",
    "                if is_zipping: \n",
    "                    zipped_name = f\"{cleanup_dir.name}_{time.strftime('%Y-%m-%d_%H%M')}.zip\" \n",
    "                    !zip -r {zipped_name} {cleanup_dir}\n",
    "                    display(f\"{zipped_name} zipped!\")\n",
    "                if is_cleaning :\n",
    "                    folder_name=cleanup_dir.name\n",
    "                    !rm -rf {cleanup_dir}\n",
    "                    cleanup_dir.mkdir(exist_ok=True)\n",
    "                    display(f\"{folder_name} cleaned!\")\n",
    "                bb.button_style = 'success' \n",
    "                bb.icon=\"check\"\n",
    "                \n",
    "            else :\n",
    "                display(\"Empty folder.\")\n",
    "                bb.button_style = 'warning' \n",
    "        \n",
    "        #listeners\n",
    "        zip_button.on_click(on_zip)\n",
    "        #clean_button.on_click(on_zip)\n",
    "        zip_clean_button.on_click(on_zip)\n",
    "    \n",
    "        #return ui\n",
    "        return HBox([\n",
    "            output_control, \n",
    "            zip_button, \n",
    "            #clean_button,\n",
    "            zip_clean_button\n",
    "        ])\n",
    "           \n",
    "   \n",
    "    #---- views ----#\n",
    "    # ressources manager\n",
    "    def view_links(self) -> None:\n",
    "        for category, urls in params.links.items():\n",
    "            if category != \"ui\":\n",
    "                display(HBox([Box(layout=Layout(width='18%')), self.get_accordion(category) ]))\n",
    "\n",
    "    # output manager \n",
    "    def show_zip_tools(self)-> None:\n",
    "        display(self.create_zip_tool())\n",
    "\n",
    "\n",
    "    \n",
    "    @property \n",
    "    def options(self) -> Any:\n",
    "        return self.links[\"ui\"][\"params\"] \n",
    "\n",
    "    @property\n",
    "    def commandline_arguments(self) -> str:\n",
    "        arg_list = [\"--enable-insecure-extension-access --disable-safe-unpickle --theme dark  --listen --port 6006\"]\n",
    "        if self.options[\"hf_token\"] != \"\":\n",
    "            arg_list.extend([\"--hf-token-out\", self.hf_token])\n",
    "        if self.options[\"update_extensions\"]  is True:\n",
    "            arg_list.append(\"--update-all-extensions\")\n",
    "        if self.options[\"save_config\"]  is True and self.options[\"config_file\"] != \"\":\n",
    "            arg_list.extend([\"--ui-settings-file\", self.options[\"config_file\"]])\n",
    "        if self.options[\"save_tasks\"]  is True :\n",
    "            arg_list.extend([\"--agent-scheduler-sqlite-file\", \"/notebooks/cached/tasks.sqlite3\"])\n",
    "        if self.options[\"extra_cmd\"] != \"\":\n",
    "            arg_list.extend([x for x in self.options[\"extra_cmd\"].split(\" \") if x != \"\"])\n",
    "        return \" \".join(arg_list)\n",
    "\n",
    "# css\n",
    "style = {\"description_width\": \"25%\"}\n",
    "layout = {\"width\": \"70%\"}\n",
    "\n",
    "label_style = HTML(\n",
    "    \"<style> .widget-label, .widget-label-basic >  span { color: white !important; font-weight: bold;}  </style>\"\n",
    ")\n",
    "\n",
    "# checkboxes\n",
    "cached_config_control = Checkbox(\n",
    "    value=False,  # Keep it false as default, presistent is good but I might change config in the future, also it's much more idiot-proof this way\n",
    "    description=\"Persistent config\",\n",
    "    layout=layout,\n",
    "    style=style,\n",
    "    tooltip='save_config',\n",
    ")\n",
    "\n",
    "cached_tasks_control = Checkbox(\n",
    "    value=False,  \n",
    "    description=\"Persistent task scheduler\",\n",
    "    layout=layout,\n",
    "    style=style,\n",
    "    tooltip='save_tasks',\n",
    ")\n",
    "\n",
    "persistent_output_control = Checkbox(\n",
    "    value=False,  # Keep it false as default, I just prefer it off by default like in old colab\n",
    "    description=\"Save outputs to notebooks\",\n",
    "    layout=layout,\n",
    "    style=style,\n",
    "    tooltip='save_output',\n",
    ")\n",
    "\n",
    "dep_size=get_directory_size(\"/storage/a1111_dependencies\")\n",
    "force_cache_control = Checkbox(\n",
    "    value=False,  # Keep it false as default, cuz why?!\n",
    "    description= f\"Force re-cache dependencies [{dep_size}]\",\n",
    "    layout=layout,\n",
    "    style=style,\n",
    "    tooltip='force_cached',\n",
    ")\n",
    "trash_size=get_directory_size('/notebooks/.Trash*')  \n",
    "clean_trash_control = Checkbox(\n",
    "    value=False,  \n",
    "    description=f\"Clean trash folder [{trash_size}]\",\n",
    "    layout=layout,\n",
    "    style=style,\n",
    "    tooltip='clean_trash',\n",
    ")\n",
    "\n",
    "cache_size=get_directory_size('/notebooks/cached')\n",
    "clean_cached_control = Checkbox(\n",
    "    value=True,  \n",
    "    description=f\"Force re-cache resources [{cache_size}] (removes tasks too)\",\n",
    "    layout=layout,\n",
    "    style=style,\n",
    "    tooltip='clean_cached',\n",
    ")\n",
    "\n",
    "extn_update_control = Checkbox(\n",
    "    value=True, # You would want newer extensions by default\n",
    "    description=\"Update all extensions on launch\",\n",
    "    layout=layout,\n",
    "    style=style,\n",
    "    tooltip='update_extensions',\n",
    ")\n",
    "\n",
    "# textboxes\n",
    "hf_token = Text(\n",
    "    value=\"\",\n",
    "    description=\"HuggingFace token (optional)\",\n",
    "    placeholder=\"hf_xyz\",\n",
    "    style=style,\n",
    "    layout=layout,\n",
    "    tooltip='hf_token',\n",
    ")\n",
    "commandline_arg_control = Text(\n",
    "    value=\"--xformers\",\n",
    "    description=\"Extra commandline arguments\",\n",
    "    style=style,\n",
    "    layout=layout,\n",
    "    tooltip='extra_cmd',\n",
    ")\n",
    "output_control = Text(\n",
    "    value=\"/notebooks/outputs/\",\n",
    "    description=\"Outputs path\",\n",
    "    placeholder=\"/notebooks/outputs/\",\n",
    "    style= {\"description_width\": \"46%\"},\n",
    "    layout=Layout(width='38%'),\n",
    "    tooltip='zip_folder',\n",
    ")\n",
    "\n",
    "# params instance with all the controls to save\n",
    "params=Params([cached_config_control,cached_tasks_control,persistent_output_control,force_cache_control,clean_cached_control,clean_trash_control,extn_update_control,output_control,hf_token,commandline_arg_control])\n",
    "params.load('links.json')\n",
    "\n",
    "# textboxes\n",
    "display(\n",
    "    hf_token,\n",
    "    commandline_arg_control,\n",
    ")\n",
    "params.show_zip_tools(),\n",
    "\n",
    "# checkboxes\n",
    "display(\n",
    "    cached_config_control,\n",
    "    cached_tasks_control,\n",
    "    persistent_output_control,\n",
    "    force_cache_control,\n",
    "    clean_cached_control,\n",
    "    clean_trash_control,\n",
    "    extn_update_control,\n",
    ")\n",
    "\n",
    "# show ressource manger\n",
    "params.view_links()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Start 🚀"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Run this when you're ready    . \n",
    "\n",
    "# Can't run without params\n",
    "class StopExecution(Exception):\n",
    "    def _render_traceback_(self):\n",
    "        pass\n",
    "try:\n",
    "    params\n",
    "except NameError:\n",
    "    display(\"Please run the configurations block\")\n",
    "    raise StopExecution\n",
    "\n",
    "# showing the future link\n",
    "display(f\"Will be ready in a few at https://tensorboard-{getenv('PAPERSPACE_FQDN')} ...\")\n",
    "\n",
    "# read controllers and save parameters\n",
    "params.update()\n",
    "params.save()\n",
    "\n",
    "#cleaning trash\n",
    "if params.options[\"clean_cached\"]:\n",
    "    display(\"Cleaning cashed files... \" )\n",
    "    !rm -rf '/notebooks/cached' \n",
    "                \n",
    "#cleaning trash\n",
    "if params.options[\"clean_trash\"]:\n",
    "    trash_size=get_directory_size('/notebooks/.Trash*')\n",
    "    if trash_size!=\"0G\" and  trash_size!=\"0B\" :\n",
    "        display(\"Cleaning trash files... \" + trash_size)\n",
    "        directories = glob.glob('/notebooks/.Trash*')\n",
    "        for directory in directories:\n",
    "            try:\n",
    "                check_output(['rm', '-rf', directory])\n",
    "                display(f\"Deleted: {directory}\")  \n",
    "            except CalledProcessError as e:\n",
    "                display(f\"Error deleting {directory}: {e}\")    \n",
    "            \n",
    "# paperspace-provided directories\n",
    "notebook_dir = Path(\"/notebooks\")\n",
    "storage_dir = Path(\"/storage\")\n",
    "\n",
    "# local-only semi-persistent storage\n",
    "content_dir = Path(\"/content\")\n",
    "\n",
    "# # storage directory\n",
    "# storage_symlink = notebook_dir.joinpath(\"storage\")\n",
    "# if not all(\n",
    "#     (\n",
    "#         storage_symlink.is_symlink(),\n",
    "#         storage_symlink.resolve() == storage_dir,\n",
    "#     )\n",
    "# ):\n",
    "#     # recreate the symlink\n",
    "#     storage_symlink.unlink(missing_ok=True)\n",
    "#     storage_symlink.symlink_to(storage_dir, target_is_directory=True)\n",
    "\n",
    "\n",
    "# symlink the directory for easier access\n",
    "content_dir.mkdir(exist_ok=True)\n",
    "content_symlink = notebook_dir.joinpath(\"content\")\n",
    "if not all(\n",
    "    (\n",
    "        content_symlink.is_symlink(),\n",
    "        content_symlink.resolve() == content_dir,\n",
    "    )\n",
    "):\n",
    "    # recreate the symlink\n",
    "    content_symlink.unlink(missing_ok=True)\n",
    "    content_symlink.symlink_to(content_dir, target_is_directory=True)\n",
    "\n",
    "# where we put deps\n",
    "dep_cache_dir = storage_dir.joinpath(\"a1111_dependencies\")\n",
    "\n",
    "# where models will go\n",
    "sdw_dir = content_dir.joinpath(\"sdw\")\n",
    "models_dir = sdw_dir.joinpath(\"models\")\n",
    "extensions_dir = sdw_dir.joinpath(\"extensions\")\n",
    "output_dir = notebook_dir.joinpath(\"outputs/\") \n",
    "cached_dir = notebook_dir.joinpath(\"cached\")\n",
    "if not cached_dir.exists():\n",
    "    cached_dir.mkdir(exist_ok=True)\n",
    "    \n",
    "# download repo if needed\n",
    "config_json = content_dir.joinpath(\"sdw/config.json\")\n",
    "cached_config_file = notebook_dir.joinpath(\"config.json\")\n",
    "\n",
    "if not config_json.exists() :\n",
    "    # repo isn't here so LET'S GOOOOOOOOOO\n",
    "    aria2_task = \"\\n\".join(\n",
    "        [\n",
    "            \"https://huggingface.co/NoCrypt/fast-repo/resolve/main/repo.tar.lz4\",\n",
    "            \"\\tout=repo.tar.lz4\",\n",
    "            \"https://huggingface.co/NoCrypt/fast-repo/resolve/main/cache.tar.lz4\",\n",
    "            \"\\tout=cache.tar.lz4\",\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    with NamedTemporaryFile(\"w\") as aria2_file:\n",
    "        aria2_file.write(aria2_task)\n",
    "        aria2_file.flush()\n",
    "        display(\"Downloading repo and cache. This may take a while...\")\n",
    "        check_call(\n",
    "            [\"aria2c\", f\"--input-file={aria2_file.name}\"] + aria2_opts,\n",
    "            cwd=content_dir,\n",
    "        )\n",
    "\n",
    "    display(\"Extracting repo...\")\n",
    "    check_call([\"tar\", \"-xI\", \"lz4\", \"-f\", \"repo.tar.lz4\", \"--directory=/\"], cwd=content_dir)\n",
    "    content_dir.joinpath(\"repo.tar.lz4\").unlink()\n",
    "\n",
    "    display(\"Extracting cache...\")\n",
    "    check_call([\"tar\", \"-xI\", \"lz4\", \"-f\", \"cache.tar.lz4\", \"--directory=/\"], cwd=content_dir)\n",
    "    content_dir.joinpath(\"cache.tar.lz4\").unlink()\n",
    "\n",
    "\n",
    "# deal with persistent config \n",
    "if params.options[\"save_config\"] :\n",
    "    if not cached_config_file.exists():\n",
    "        !cp {config_json} {notebook_dir}\n",
    "    config_json=cached_config_file\n",
    "params.options[\"config_file\"]= str(config_json)\n",
    "\n",
    "# deal with persistent output \n",
    "if params.options[\"save_output\"]:\n",
    "    if not output_dir.exists():\n",
    "        output_dir.mkdir(exist_ok=True)\n",
    "    !sed -i 's@\"outdir_txt2img_samples\": \"outputs/txt2img-images\"@\"outdir_txt2img_samples\": \"{output_dir}/txt2img-images\"@' {config_json}\n",
    "    !sed -i 's@\"outdir_img2img_samples\": \"outputs/img2img-images\"@\"outdir_img2img_samples\": \"{output_dir}/img2img-images\"@' {config_json}\n",
    "    !sed -i 's@\"outdir_extras_samples\": \"outputs/extras-images\"@\"outdir_extras_samples\": \"{output_dir}/extras-images\"@' {config_json}\n",
    "    !sed -i 's@\"outdir_txt2img_grids\": \"outputs/txt2img-grids\"@\"outdir_txt2img_grids\": \"{output_dir}/txt2img-grids\"@' {config_json}\n",
    "    !sed -i 's@\"outdir_img2img_grids\": \"outputs/img2img-grids\"@\"outdir_img2img_grids\": \"{output_dir}/img2img-grids\"@' {config_json}\n",
    "    !sed -i 's@\"outdir_save\": \"log/images\"@\"outdir_save\": \"{output_dir}/log/images\"@' {config_json}\n",
    "\n",
    "# remove unwanted extensions from nocrypt packaging\n",
    "pixelization_dir = extensions_dir.joinpath(\"pixelization\")\n",
    "photope_dir = extensions_dir.joinpath(\"sd-photopea-embed-lazy\")\n",
    "!rm -rf {pixelization_dir}\n",
    "!rm -rf {photope_dir}\n",
    "    \n",
    "#Adds wildcards\n",
    "wildcards_dir_src = notebook_dir.joinpath(\"wildcards\")\n",
    "wildcards_dir_dest = extensions_dir.joinpath(\"Umi-AI-debloat/\")\n",
    "if not wildcards_dir_src.exists():\n",
    "    wildcards_dir_src.mkdir(exist_ok=True)\n",
    "if wildcards_dir_dest.exists():\n",
    "    !cp -r {wildcards_dir_src} {wildcards_dir_dest}\n",
    "\n",
    "# flush dependency cache if forced\n",
    "if params.options[\"force_cached\"] :\n",
    "    if dep_cache_dir.exists() :\n",
    "        display(\"Forced removal of dependencies...\")\n",
    "        shutil.rmtree(dep_cache_dir)\n",
    "\n",
    "current_dep_dir=Path(get_python_package_path())\n",
    "\n",
    "# set up dependency cache\n",
    "if dep_cache_dir.exists() :\n",
    "    if current_dep_dir.exists():\n",
    "        display(\"Copying dependencies from cache...\")\n",
    "\n",
    "        # recursively copy back files from the cached folder to the current one \n",
    "        shutil.copytree(\n",
    "            dep_cache_dir,\n",
    "            current_dep_dir,\n",
    "            symlinks=True,\n",
    "            dirs_exist_ok=True,\n",
    "        )\n",
    "else :\n",
    "    dep_cache_dir.mkdir(exist_ok=True)\n",
    "    start_time = time.time()\n",
    "    check_call(\n",
    "        [\"bash\", \"./webui.sh\", \"-f\", params.commandline_arguments, \"--exit\"],\n",
    "        cwd=content_dir.joinpath(\"sdw\"),\n",
    "    )\n",
    "    end_time = time.time()\n",
    "    # updated no need to update later on\n",
    "    params.options[\"update_extensions\"]=False\n",
    "    copy_recent_files(current_dep_dir, dep_cache_dir, round((end_time - start_time) / 60)+1)\n",
    "\n",
    "# func to download resources\n",
    "def acquire_checkpoints(checkpoint_urls: list[str], fldr: str = \"Stable-diffusion\", is_cached: bool = False) -> None:\n",
    "\n",
    "    sdw_fldr=paths.get(fldr, fldr)\n",
    "    dst=sdw_dir.joinpath(sdw_fldr)\n",
    "    real_dst=sdw_dir.joinpath(sdw_fldr)\n",
    "\n",
    "    # cleaning symlinks and cache uncached\n",
    "    if(is_cached) :\n",
    "        #cleaning symbolic links\n",
    "        sym_dst=sdw_dir.joinpath(paths.get(category, category))\n",
    "        if sym_dst.exists() :\n",
    "            delete_symlinks(sym_dst)\n",
    "        #cleaning cached\n",
    "        dst=cached_dir.joinpath(fldr)\n",
    "        if not dst.exists():\n",
    "            dst.mkdir(exist_ok=True)\n",
    "        else :\n",
    "            for item in dst.iterdir(): \n",
    "                if item.name.startswith(\".\") :\n",
    "                    continue\n",
    "                url=params.link_from_name(item.name)\n",
    "                enabled, keep=params.get_link_status(url)\n",
    "                if not keep and not enabled:\n",
    "                    #delete from cache\n",
    "                    item.unlink()\n",
    "                    \n",
    "                elif not keep and enabled:\n",
    "                    # move away from cache\n",
    "                    moved_dst=real_dst.joinpath(item.name)\n",
    "                    if not moved_dst.exists() :\n",
    "                         shutil.move(item, real_dst)\n",
    "                    else :\n",
    "                        item.unlink()               \n",
    "        if real_dst.exists():\n",
    "            for item in real_dst.iterdir(): \n",
    "                if item.name.startswith(\".\") :\n",
    "                    continue\n",
    "                url=params.link_from_name(item.name)\n",
    "                enabled, keep=params.get_link_status(url)\n",
    "                if  keep and enabled:\n",
    "                    # move away from cache\n",
    "                    moved_dst=dst.joinpath(item.name)\n",
    "                    if not moved_dst.exists() :\n",
    "                        shutil.move(item,dst)\n",
    "                    else :\n",
    "                        item.unlink()        \n",
    "                                  \n",
    "    if len(checkpoint_urls) == 0 :\n",
    "        return  \n",
    "    aria2_task_lines = []\n",
    "    for url in checkpoint_urls:\n",
    "        filename=params.name_from_link(url)\n",
    "        if fldr==\"extensions\" :\n",
    "            if filename == \"\" :\n",
    "                filename=urlparse(url).path.split('/')[-1]\n",
    "                filename=filename.replace('.git','')\n",
    "            if not dst.joinpath(filename).exists():\n",
    "                !git clone --quiet {url} {dst}/{filename} \n",
    "            else:\n",
    "                !cd {dst}/{filename} && git fetch --quiet > /dev/null 2>&1 && git merge --quiet\n",
    "        if 'drive.google' in url:\n",
    "            if not dst.joinpath(filename).exists():\n",
    "                downloaded_file_path = gdown.download(url, str(dst)+\"/\", quiet=True, fuzzy=True) \n",
    "                filename = Path(downloaded_file_path).name\n",
    "                params.set_link_name(url,filename)\n",
    "        elif filename==\"\" :\n",
    "            filename= wget_name(url,dst)\n",
    "            if filename :\n",
    "                params.set_link_name(url,filename)\n",
    "        elif \"login\" not in filename:\n",
    "            if not dst.joinpath(filename).exists():\n",
    "                aria2_task_lines.extend([url,f\"\\tout={filename}\"])\n",
    "            params.set_link_name(url,filename)\n",
    "          \n",
    "    if len(aria2_task_lines) > 0:\n",
    "        with NamedTemporaryFile(\"w\") as aria2_file:\n",
    "            aria2_file.write(\"\\n\".join(aria2_task_lines))\n",
    "            aria2_file.flush()\n",
    "            check_call([\"aria2c\", f\"--input-file={aria2_file.name}\"] + aria2_opts,cwd=dst,)\n",
    "            \n",
    "    if is_cached :\n",
    "        for file in dst.iterdir():\n",
    "            # hidden files\n",
    "            if file.name.startswith('.') or file.name.endswith('.aria2') :  \n",
    "                continue\n",
    "\n",
    "            #extensions\n",
    "            if fldr ==\"extensions\":\n",
    "                extension_symlink = Path(real_dst / file.name) \n",
    "                if not all(\n",
    "                    (\n",
    "                        extension_symlink.is_symlink(),\n",
    "                        extension_symlink.resolve() == file,\n",
    "                    )\n",
    "                ):\n",
    "                    # recreate the symlink\n",
    "                    extension_symlink.unlink(missing_ok=True)\n",
    "                    extension_symlink.symlink_to(file, target_is_directory=True)\n",
    "                continue\n",
    "\n",
    "            #other resources\n",
    "            original_file = real_dst / file.name\n",
    "            symlink_path = file\n",
    "            if original_file.exists() or original_file.is_symlink():\n",
    "                #display(f\"Removing existing file or symlink: {original_file}\")\n",
    "                original_file.unlink()\n",
    "            try:\n",
    "                original_file.symlink_to(symlink_path)\n",
    "                #display(f\"Created symbolic link: {original_file} -> {symlink_path}\")\n",
    "            except FileExistsError as e:\n",
    "                #display(f\"Failed to create symbolic link due to existing file: {e}\")\n",
    "                pass\n",
    "  \n",
    "\n",
    "# download, cache and link ressources \n",
    "for category, urls in params.links.items():\n",
    "    if category ==\"ui\" :\n",
    "        continue\n",
    "  \n",
    "    active_urls= []\n",
    "    persistent_urls= []\n",
    "        \n",
    "    if len(urls) >0 :\n",
    "        for url in urls:\n",
    "            enabled, keep = params.get_link_status(url)\n",
    "            if enabled :\n",
    "                if not keep :\n",
    "                    if url not in active_urls:\n",
    "                        active_urls.append(url)\n",
    "                    if url in persistent_urls:\n",
    "                        persistent_urls.remove(url)\n",
    "                else :\n",
    "                    if url not in persistent_urls:\n",
    "                        persistent_urls.append(url)\n",
    "                    if url in active_urls:\n",
    "                        active_urls.remove(url)\n",
    "            else:\n",
    "                if url in active_urls:\n",
    "                        active_urls.remove(url)\n",
    "                if url in persistent_urls:\n",
    "                        persistent_urls.remove(url)\n",
    "\n",
    "    \n",
    "    display(f\"[{category}] downloading... {(len(persistent_urls)+len(active_urls))} resources\")\n",
    "    acquire_checkpoints(persistent_urls, category, True) \n",
    "    acquire_checkpoints(active_urls,category)\n",
    "    \n",
    "params.save()\n",
    "\n",
    "# start. the thing\n",
    "environ[\"COMMANDLINE_ARGS\"] = params.commandline_arguments\n",
    "environ[\"REQS_FILE\"] = \"requirements_versions.txt\"\n",
    "\n",
    "display(\"Starting 'tensorboard'...\")\n",
    "display(f\"Access via this URL: https://tensorboard-{getenv('PAPERSPACE_FQDN')}\")\n",
    "\n",
    "webui_user = f\"\"\"#!/bin/bash\n",
    "export COMMANDLINE_ARGS='{params.commandline_arguments}'\n",
    "venv_dir='-'\n",
    "\"\"\"\n",
    "content_dir.joinpath(\"sdw/webui-user.sh\").write_text(webui_user)\n",
    "check_call([\"/usr/bin/env\", \"bash\", \"./webui.sh\", \"-f\"], cwd=content_dir.joinpath(\"sdw\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# code to update webui\n",
    "!cd {sdw_dir} && git pull"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}