comfyui_colab_google_drive.ipynb

comfyui_colab_google_drive.ipynb

import os
import tarfile
import time
from google.colab import drive

start_time = time.time()

RE_BACKUP = False  #@param {type:"boolean"}
MASTER_RESET = False  #@param {type:"boolean"}


##### Define paths
LOCAL_WORKSPACE = "/content"
LOCAL_VIRTUAL_ENV_NAME = "colab_env"

LOCAL_VIRTUSAL_ENV_PATH = f"{LOCAL_WORKSPACE}/{LOCAL_VIRTUAL_ENV_NAME}"
REMOTE_DRIVE_MOUNT_PATH = f"{LOCAL_WORKSPACE}/drive"

REMOTE_DRIVE_PATH = f"{REMOTE_DRIVE_MOUNT_PATH}/MyDrive"
REMOTE_DRIVE_COMFYUI_WORKSPACE = f"{REMOTE_DRIVE_PATH}/ComfyUI"

REMOTE_VIRTUAL_ENV_BACKUP_ARCHIVE = f"{REMOTE_DRIVE_PATH}/{LOCAL_VIRTUAL_ENV_NAME}.tar.gz"

PYTHON_EXEC = f"{LOCAL_VIRTUSAL_ENV_PATH}/bin/python"
os.environ["COMFYUI_PATH"] = REMOTE_DRIVE_COMFYUI_WORKSPACE

print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
if(not os.path.exists(REMOTE_DRIVE_PATH)):
  print("Google Drive is not mounted, Please mount google drive first")
  drive.mount(REMOTE_DRIVE_MOUNT_PATH)
  print("Google Drive Mount Complete")
else:
  print("Google Drive Mounted, no action needed")


print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
if(os.path.exists(REMOTE_VIRTUAL_ENV_BACKUP_ARCHIVE) and MASTER_RESET):
  print("STEP :  Deleting backup gogole drive")
  !rm $REMOTE_VIRTUAL_ENV_BACKUP_ARCHIVE

##### Change to working directory
print("STEP :  Change to working directory {}")
%cd $LOCAL_WORKSPACE

##### Restore virtual environment if it exists in Drive
print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
print("STEP : System Setup")
if os.path.exists(REMOTE_VIRTUAL_ENV_BACKUP_ARCHIVE):
    print("STEP : System Backup found, Downloading")
    !tar --checkpoint=100000 --checkpoint-action=echo="%T files extracted" -xzf $REMOTE_VIRTUAL_ENV_BACKUP_ARCHIVE -C $LOCAL_WORKSPACE/
    print("STEP : Download Complete")
    %cd $REMOTE_DRIVE_COMFYUI_WORKSPACE
else:
    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print("STEP : Boostraping the enviroment for the first time")
    
    !pip install virtualenv
    !virtualenv $LOCAL_VIRTUSAL_ENV_PATH
    print("STEP : Boostraping the enviroment complete")

    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print("STEP : Activiating Virtual enviroment")
    !bash -c "source $LOCAL_VIRTUSAL_ENV_PATH/bin/activate && echo Virtual Environment Activated"

    ##### Clone ComfyUI if not already present
    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    if not os.path.exists(REMOTE_DRIVE_COMFYUI_WORKSPACE):
      print("STEP : Downloading ComfyUI")
      !git clone https://github.com/comfyanonymous/ComfyUI $REMOTE_DRIVE_COMFYUI_WORKSPACE


    ##### Clone ComfyUI-Manager
    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print("STEP : Downloading ComfyUI-Manager")
    %cd $REMOTE_DRIVE_COMFYUI_WORKSPACE/custom_nodes
    if not os.path.exists("comfyui-manager"):
        !git clone https://github.com/ltdrdata/ComfyUI-Manager $REMOTE_DRIVE_COMFYUI_WORKSPACE/custom_nodes/comfyui-manager


    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print("STEP : Installing Dependecies")
    %cd $REMOTE_DRIVE_COMFYUI_WORKSPACE
    !$PYTHON_EXEC -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
    !$PYTHON_EXEC -m pip install -r requirements.txt
    !$PYTHON_EXEC -m pip install -r custom_nodes/comfyui-manager/requirements.txt


    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print("STEP : Running Comfy UI Manager")
    !$PYTHON_EXEC custom_nodes/comfyui-manager/cm-cli.py install comfy-image-saver
    !$PYTHON_EXEC custom_nodes/comfyui-manager/cm-cli.py restore-dependencies
    print("STEP : Boostraping Complete")

####### Save environment after installing dependencies (only needed once)
if ((not os.path.exists(REMOTE_VIRTUAL_ENV_BACKUP_ARCHIVE)) or RE_BACKUP):
    print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
    print("STEP : Creating a backup and storing in google drive")
    !tar --checkpoint=100000 --checkpoint-action=echo="Processed %T files" -czf $REMOTE_VIRTUAL_ENV_BACKUP_ARCHIVE -C $LOCAL_WORKSPACE $LOCAL_VIRTUAL_ENV_NAME

print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
print("STEP : System Setup Complete")
print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
end_time = time.time()
elapsed_time = end_time - start_time
print(f"Elapsed Time {elapsed_time:.2f} seconds")
print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")