Omost.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Omost - LLM + Canvas (Regional Prompting) + Stable Diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title ## Clone the repo and install requirements- Omost { display-mode: \"form\" }\n",
    "\n",
    "!git clone https://github.com/lllyasviel/Omost.git\n",
    "%cd Omost\n",
    "#!conda create -n omost python=3.10\n",
    "#!conda activate omost\n",
    "#!pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121\n",
    "!pip install -r requirements.txt\n",
    "\n",
    "#!python gradio_app.py\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title ## Edit the files { display-mode: \"form\" }\n",
    "\n",
    "# function to edit the files\n",
    "def Edit_file_code(goal, target_file, target_code, new_code, lines_offset, pop_x_lines):\n",
    "\n",
    "    print(f'''-------------------------------------\n",
    "Purpose: {goal}\n",
    "target file: {target_file}\n",
    "target code:\\n{target_code}\n",
    "new code:\\n{new_code}\n",
    "lines offset: {lines_offset}\n",
    "{pop_x_lines} lines to pop''')\n",
    "\n",
    "    # check if the edit was already done, exit if so\n",
    "    with open(target_file, 'r') as file:\n",
    "        data = file.read()\n",
    "    if (new_code in data):\n",
    "        print (f'''The new code was found in target file: {target_file}. The edit was already done. Exiting.''')\n",
    "        return\n",
    "\n",
    "    # read the file\n",
    "    with open(target_file, 'r') as file:\n",
    "        lines = file.readlines()\n",
    "\n",
    "    found_line_number = -1\n",
    "\n",
    "    # find the target line \"target_code\"\n",
    "    for i, line in enumerate(lines):\n",
    "        if target_code in line:\n",
    "            found_line_number = i\n",
    "            break\n",
    "\n",
    "    # it should always find it, but whatever, insert \"new_code\" to the data\n",
    "    if found_line_number != -1:\n",
    "        if (pop_x_lines):\n",
    "            print (\"haa\")\n",
    "            for i in range(pop_x_lines):\n",
    "                lines.pop(found_line_number)\n",
    "        lines.insert(found_line_number + lines_offset, new_code)\n",
    "        print (f'''New code added to target file: {target_file}. File edited.''')\n",
    "    else:\n",
    "        print (f'''Could not find target code: {target_code}.''')\n",
    "\n",
    "    # rewrite target_file\n",
    "    with open(target_file, 'w') as file:\n",
    "        file.writelines(lines)\n",
    "\n",
    "\n",
    "# functions for every edit/changes to be made\n",
    "\n",
    "# to fix an error, can't remember what it was anymore, it works\n",
    "def edit_memory_management_py_first_change():\n",
    "    # variables\n",
    "    goal = \"First change\"\n",
    "    target_file = \"/content/Omost/lib_omost/memory_management.py\"\n",
    "    target_code = '''    if not high_vram:'''\n",
    "    new_code = r'''    if not high_vram:\n",
    "        for m in models_to_unload:\n",
    "            if m.__class__.__name__ != \"LlamaForCausalLM\":\n",
    "                with movable_bnb_model(m):\n",
    "                    m.to(cpu)\n",
    "                    print('Unload to CPU:', m.__class__.__name__)\n",
    "        models_in_gpu = models_to_remain\n",
    "'''\n",
    "    lines_offset = 0\n",
    "    pop_x_lines = 6\n",
    "    # call the edit file function\n",
    "    Edit_file_code(goal, target_file, target_code, new_code, lines_offset, pop_x_lines)\n",
    "\n",
    "# to fix an error, something about 2 devices, this sets it to gpu (I think? can't remember)\n",
    "def edit_gradio_app_py_first_change():\n",
    "    # variables\n",
    "    goal = \"First change\"\n",
    "    target_file = \"/content/Omost/gradio_app.py\"\n",
    "    target_code = '''import os'''\n",
    "    new_code = r'''import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "'''\n",
    "    lines_offset = 0\n",
    "    pop_x_lines = 1\n",
    "    # call the edit file function\n",
    "    Edit_file_code(goal, target_file, target_code, new_code, lines_offset, pop_x_lines)\n",
    "\n",
    "# to get a gradio public/share link\n",
    "def edit_gradio_app_py_share_link():\n",
    "    # variables\n",
    "    goal = \"First change\"\n",
    "    target_file = \"/content/Omost/gradio_app.py\"\n",
    "    target_code = '''    demo.queue().launch(inbrowser=True, server_name='0.0.0.0')'''\n",
    "    new_code = r'''    demo.queue().launch(inbrowser=True, share=True, server_name='0.0.0.0')\n",
    "'''\n",
    "    lines_offset = 0\n",
    "    pop_x_lines = 1\n",
    "    # call the edit file function\n",
    "    Edit_file_code(goal, target_file, target_code, new_code, lines_offset, pop_x_lines)\n",
    "\n",
    "\n",
    "# to fix an error, This was supposed to unload all models to prevent OOM, doesn't work, kills everything instead when generating the image\n",
    "def edit_memory_management_py_second_change_mine():\n",
    "    # variables\n",
    "    goal = \"Second change\"\n",
    "    target_file = \"/content/Omost/lib_omost/memory_management.py\"\n",
    "    target_code = '''def unload_all_models(extra_models=None):'''\n",
    "    new_code = r'''def unload_all_models(extra_models=None):\n",
    "    global models_in_gpu\n",
    "\n",
    "    if extra_models is None:\n",
    "        extra_models = []\n",
    "\n",
    "    if not isinstance(extra_models, (tuple, list)):\n",
    "        extra_models = [extra_models]\n",
    "\n",
    "    models_in_gpu = list(set(models_in_gpu + extra_models))\n",
    "\n",
    "    for m in models_in_gpu:\n",
    "        with movable_bnb_model(m):\n",
    "            m.to(cpu)\n",
    "            print('Unload from GPU:', m.__class__.__name__)\n",
    "            # Unload from CPU\n",
    "            del m\n",
    "    models_in_gpu = []\n",
    "    torch.cuda.empty_cache()\n",
    "    \n",
    "    return load_models_to_gpu([])\n",
    "'''\n",
    "    lines_offset = 0\n",
    "    pop_x_lines = 12\n",
    "    # call the edit file function\n",
    "    Edit_file_code(goal, target_file, target_code, new_code, lines_offset, pop_x_lines)\n",
    "\n",
    "\n",
    "\n",
    "# call the changes\n",
    "# These changes fix the first 2 errors, but it overloads the memory, both ram and vram\n",
    "edit_memory_management_py_first_change()\n",
    "edit_gradio_app_py_first_change()\n",
    "# gradio public link\n",
    "edit_gradio_app_py_share_link()\n",
    "\n",
    "# This is supposed to unload all models to prevent OOM, doesn't work, kills the app\n",
    "# suggested by you.com, goes to show what it knows\n",
    "#edit_memory_management_py_second_change_mine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python gradio_app.py"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}