Madiator2011/madiator-docker-runpod
1
0
Fork 0
mirror of https://github.com/kodxana/madiator-docker-runpod.git synced 2024-11-22 02:40:12 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #3 from lutzapps/shared_models

Browse source 
Initial Changes for localhost support and extended shared_models module
This commit is contained in:
Madiator2011 2024-10-30 15:44:34 +01:00 committed by GitHub
commit 458f331f13
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
33 changed files with 2774 additions and 344 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
official-templates/better-ai-launcher/.dockerignore Normal file
View file

@ -0,0 +1,31 @@
**/.DS_Store
**/.dockerenv
**/filebrowser.db
**/docker-bake.hcl
**/__pycache__
**/.venv
**/.classpath
**/.dockerignore
**/.env
**/.git
**/.gitignore
**/.project
**/.settings
**/.toolstarget
**/.vs
**/.vscode
**/*.*proj.user
**/*.dbmdl
**/*.jfm
**/bin
**/charts
**/docker-compose*
**/compose*
**/Dockerfile*
**/node_modules
**/npm-debug.log
**/obj
**/secrets.dev.yaml
**/values.dev.yaml
LICENSE
README.md

5
official-templates/better-ai-launcher/.gitignore vendored Normal file
View file

@ -0,0 +1,5 @@
**/.DS_Store
**/.dockerenv
**/__pycache__
**/.venv
**/.env

28
official-templates/better-ai-launcher/.vscode/launch.json vendored Normal file
View file

@ -0,0 +1,28 @@
{
"configurations": [
{
"name": "Docker: Python - Flask",
"type": "docker",
"request": "launch",
"preLaunchTask": "docker-run: debug",
"python": {
"pathMappings": [
{
"localRoot": "${workspaceFolder}",
"remoteRoot": "/app"
}
],
"projectType": "flask"
}
},
{ // this whole section was added manually to add "justMyCode": false
"name": "Python: Debug Tests",
"type": "debugpy",
"request": "launch",
"program": "${file}",
"purpose": ["debug-test"],
"console": "integratedTerminal",
"justMyCode": false
}
]
}

72
official-templates/better-ai-launcher/.vscode/tasks.json vendored Normal file
View file

@ -0,0 +1,72 @@
{
"version": "2.0.0",
"tasks": [
{
"type": "docker-build",
"label": "docker-build",
"platform": "python",
"dockerBuild": {
"tag": "madiator2011/better-launcher:dev",
"dockerfile": "${workspaceFolder}/Dockerfile",
"context": "${workspaceFolder}",
"pull": true
}
},
{
"type": "docker-run",
"label": "docker-run: debug",
"dependsOn": [
"docker-build"
],
"dockerRun": {
"containerName": "madiator2011-better-launcher", // no "/" allowed here for container name
"image": "madiator2011/better-launcher:dev",
"envFiles": ["${workspaceFolder}/.env"], // pass additional env-vars (hf_token, civitai token, ssh public-key) from ".env" file to container
"env": { // this ENV vars go into the docker container to support local debugging
"LOCAL_DEBUG": "True", // change app to localhost Urls and local Websockets (unsecured)
"FLASK_APP": "app/app.py",
"FLASK_ENV": "development", // changed from "production"
"GEVENT_SUPPORT": "True" // gevent monkey-patching is being used, enable gevent support in the debugger
// "FLASK_DEBUG": "0" // "1" allows debugging in Chrome, but then VSCode debugger not works
},
"volumes": [
{
"containerPath": "/app",
"localPath": "${workspaceFolder}" // the "/app" folder (and sub-folders) will be mapped locally for debugging and hot-reload
},
{
"containerPath": "/workspace",
// TODO: create the below folder before you run!
"localPath": "${userHome}/Projects/Docker/Madiator/workspace"
}
],
"ports": [
{
"containerPort": 7222, // main Flask app port "App-Manager"
"hostPort": 7222
},
{
"containerPort": 8181, // File-Browser
"hostPort": 8181
},
{
"containerPort": 7777, // VSCode-Server
"hostPort": 7777
}
]
},
"python": {
"args": [
"run",
// "--no-debugger", // disabled to support VSCode debugger
// "--no-reload", // disabled to support hot-reload
"--host",
"0.0.0.0",
"--port",
"7222"
],
"module": "flask"
}
}
]
}

63
official-templates/better-ai-launcher/Dockerfile
View file

@ -1,13 +1,20 @@
# Use the specified base image
FROM madiator2011/better-base:cuda12.1 as base
# lutzapps - use uppercase "AS"
FROM madiator2011/better-base:cuda12.1 AS base
# lutzapps - prepare for local developement and debugging
# needed to change the ORDER of "apt-get commands" and move the "update-alternatives" for python3
# AFTER the "apt-get remove -y python3.10" cmd, OTHERWISE the symlink to python3
# is broken in the image and the VSCode debugger could not exec "python3" as CMD overwrite
# Install Python 3.11, set it as default, and remove Python 3.10
RUN apt-get update && \
apt-get install -y python3.11 python3.11-venv python3.11-dev python3.11-distutils aria2 git \
pv git rsync zstd libtcmalloc-minimal4 bc nginx ffmpeg && \
apt-get remove -y python3.10 python3.10-minimal libpython3.10-minimal libpython3.10-stdlib && \
update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1 && \
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 && \
apt-get remove -y python3.10 python3.10-minimal libpython3.10-minimal libpython3.10-stdlib && \
apt-get autoremove -y && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
@ -23,46 +30,64 @@ WORKDIR /app
# Copy the requirements file
COPY requirements.txt .
# Install the Python dependencies
# Install the Python dependencies (as "managed pip")
RUN python3.11 -mpip install --no-cache-dir -r requirements.txt
# Copy the application code
COPY . .
# lutzapps - only copy the "app" folder and not the root (".") to avoid cluttering the docker container with src-files
COPY app .
# Install File Browser
RUN curl -fsSL https://raw.githubusercontent.com/filebrowser/get/master/get.sh | bash
# Set environment variables for production
# Set environment variables for developent/production
# overwrite this ENV in "tasks.json" docker run "env" section or overwrite in your ".env" file to "development"
ENV FLASK_ENV=production
# gevent monkey-patching is being used, enable gevent support in the debugger with GEVENT_SUPPORT=True
# add this ENV in "tasks.json" docker run "env" section or populate in your ".env" file
# ENV GEVENT_SUPPORT=True
# lutzapps - keep Python from generating .pyc files in the container
ENV PYTHONDONTWRITEBYTECODE=1
# Turns off buffering for easier container logging
ENV PYTHONUNBUFFERED=1
ENV APP_PATH=/app/app.py
# Expose the port Nginx will listen on
EXPOSE 7222
# lutzapps - moved the 4 static assets (3x PNG, 1x MP3) into the /app/static" folder for cleaner view
# lutzapps - added a "app/tests" folder with script and testdata and readme file
# lutzapps - grouped NGINX files in a sub-folder for cleaner view
# Copy the README.md
COPY README.md /usr/share/nginx/html/README.md
COPY nginx/README.md /usr/share/nginx/html/README.md
# NGINX configuration
COPY nginx.conf /etc/nginx/nginx.conf
COPY readme.html /usr/share/nginx/html/readme.html
# Create a directory for static files
RUN mkdir -p /app/static
# Copy the Poddy animation files to the static directory
COPY poddy.png /app/static/poddy.png
COPY mushroom.png /app/static/mushroom.png
COPY snake.png /app/static/snake.png
COPY poddy-song.mp3 /app/static/poddy-song.mp3
COPY nginx/nginx.conf /etc/nginx/nginx.conf
COPY nginx/readme.html /usr/share/nginx/html/readme.html
# Copy all necessary scripts
COPY --from=scripts start.sh /
COPY pre_start.sh /pre_start.sh
RUN chmod +x /pre_start.sh /start.sh
# --from=scripts is defined as a "shared" location in "docker-bake.hcl" in the "contexts" dict:
# scripts = "../../container-template"
# the local "start.sh" is (intentionally) empty
# to build all from *one* location, copy "start.sh" here into the project workspace folder first
# cp ../../container-template/scripts/start.sh start.sh
#COPY start.sh /
COPY pre_start.sh /
# lutzapps - add execution flags to added "/app/tests/populate_testdata.sh"
RUN chmod +x /pre_start.sh /start.sh /app/tests/populate_testdata.sh
# Copy the download_venv.sh script and make it executable
COPY download_venv.sh /app/download_venv.sh
RUN chmod +x /app/download_venv.sh
# CMD
# During debugging, this entry point will be overridden as "python3".
# For more information, please refer to https://aka.ms/vscode-docker-python-debug
CMD ["/start.sh"]

62
official-templates/better-ai-launcher/README-Development.txt Normal file
View file

@ -0,0 +1,62 @@
For local development of this image, you can run the image with 3 options:
- Docker Compose (production or Development)
- VS Code (F5 Debugging)
PREPARE ENV variables:
Both docker-compose YML files and the VS-Code "tasks.json" configuration file
pass a ".env" file into the container, to set User-specific ENV variables.
Edit the supplied "env.txt" template file with your values, and then rename it to ".env".
The ".env" file is in the ".gitignore" file to avoid unwanted secret-sharing with GitHub!
To build and run the image with DOCKER COMPOSE:
To run in "production":
Use the command "docker compose up":
That runs the container without debugger, but enables localhost and a workspace bind.
It uses the default "docker-compose.yml", which should be ADJUSTED to your workspace bind location.
This YML file binds all application ports (7222, 8181, 7777) all browsable from localhost,
and your SSH public key will be configured and be usable!
It uses 1 docker bind mount:
- "/workspace" can be bound to any location on your local machine (ADJUST the location in the YML file)
This option uses the default entry CMD "start.sh" as defined in the "Dockerfile".
To run in "development":
Use the command "docker compose -f docker-compose.debug.yml"
That runs the container with a python debugger (debugpy) attached, enables localhost browsing,
and binds the workspace from a local folder of your choice.
It uses 2 docker bind mounts:
- "/app" to mirror your app into the container (supports hot-reload) and debugging against your source files
- "/workspace" can be bound to any location on your local machine (ADJUST the location in the YML file)
This second debug YML configuration is configured with the same settings that are used for
VS-Code debugging with the VS-Code Docker Extension with Run -> Debug (F5).
BUT YOU NEED TO ATTACH THE DEBUGGER YOURSELF
Use the debugpy.wait_for_client() function to block program execution until the client is attached.
See more at https://github.com/microsoft/debugpy
It is much easier to build and run the image with VS-CODE (Run -> Debug F5):
The VS-Code Docker Extension uses 2 files in the hidden folder ".vscode":
- launch.json
- tasks.json
The "tasks.json" file is the file with most configuration settings
which basically mirrors to the "docker-compose.debug.yml"
ALL these 5 mention files (env.txt, docker-compose.yml, docker-compose.debug.yml,
launch.json, tasks.json, en.txt) are heavily commented.
NOTES
Both debugging options with "docker-compose.debug.yml", or VS-Code Docker Extenstions "tasks.json"
replace the CMD entry point of the Dockerfile to "python3", instead of "start.sh" !!!
That means that the apps on port 8181 (File-Browser) and 7777 (VSCode-Server) are
NOT available during debugging, neither will your SSH public key be configured and is not usable,
as these configuration things happen in "start.sh", but they are not needed during development.
Only the NGINX server and the Flask module will be started during debugging.

100
official-templates/better-ai-launcher/README.md
View file

@ -1,7 +1,99 @@
## Build Options
# madiator-docker-runpod
RunPod Docker Containers for RunPod
To build with default options, run `docker buildx bake`, to build a specific target, run `docker buildx bake <target>`.
**Better AI Launcher Container for RunPod and local develoment**
## Ports
### Build Vars ###
IMAGE_BASE=madiator2011/better-launcher
- 22/tcp (SSH)
IMAGE_TAG=dev
### Github: ###
https://github.com/kodxana/madiator-docker-runpod
### ENV Vars ###
These ENV vars go into the docker container to support local debugging:
see also explanantion in ".vscode/tasks.json" or "docker-compose.debug.yml"
LOCAL_DEBUG=True
change app to localhost Urls and local Websockets (unsecured)
FLASK_ENV=development
changed from "production" (default)
GEVENT_SUPPORT=True
gevent monkey-patching is being used, enable gevent support in the debugger
FLASK_DEBUG=0
"1" allows debugging in Chrome, but then VSCode debugger not works
*User ENV Vars for Production:*
### APP specific Vars ###
DISABLE_PULLBACK_MODELS=False
the default is, that app model files, which are found locally (in only one app),
get automatically "pulled-back" into the '/workspace/shared_models' folder.
From there they will be re-linked back not only to their own "pulled-back" model-type folder,
but also will be linked back into all other corresponding app model-type folders.
So the "pulled-back" model is automatically shared to all installed apps.
If you NOT want this behaviour, then set DISABLE_PULLBACK_MODELS=True
### USER specific Vars and Secrets (Tokens) - TODO: adjust this for your personal settings ###
PUBLIC_KEY=ssh-ed25519 xxx...xxx usermail@domain.com
HF_TOKEN=hf_xxx...xxx
CIVITAI_API_TOKEN=xxx.xxx
**SECURITY TIP:**
These 3 security sensitive environment vars should be stored as RUNPOD **SECRETS** and referenced directly in your POD Template in the format {{ RUNPOD_SECRET_MYENVVAR }}
From https://docs.runpod.io/pods/templates/secrets
You can reference your Secret directly in the Environment Variables section of your Pod template. To reference your Secret, reference it's key appended to the "RUNPOD_SECRET_" prefix.
That mean, for this template/image, you should use these formats:
{{ RUNPOD_SECRET_PUBLIC_KEY}}
{{ RUNPOD_SECRET_HF_TOKEN }}
{{ RUNPOD_SECRET_CIVITAI_API_TOKEN }}
### Ports: ###
SSH-Port
22:22/tcp
App-Manager
7222:7222/http
VSCode-Server
7777:7777/http
File-Browser
8181:8181/http
### Apps: ###
ComfyUI
3000:3000/http
Forge (Stable-Diffiusion-WebUI-Forge)
7862:7862/http
A1111 (Stable-Diffiusion-WebUI)
7863:7863/http
*coming soon*
Kohya-ss
7864:7864/http

190
official-templates/better-ai-launcher/app.py → official-templates/better-ai-launcher/app/app.py
View file

@ -17,11 +17,47 @@ from utils.filebrowser_utils import configure_filebrowser, start_filebrowser, st
from utils.app_utils import (
run_app, update_process_status, check_app_directories, get_app_status,
force_kill_process_by_name, update_webui_user_sh, save_install_status,
get_install_status, download_and_unpack_venv, fix_custom_nodes, is_process_running, install_app, update_model_symlinks
get_install_status, download_and_unpack_venv, fix_custom_nodes, is_process_running, install_app #, update_model_symlinks
)
# lutzapps - CHANGE #1
LOCAL_DEBUG = os.environ.get('LOCAL_DEBUG', 'False') # support local browsing for development/debugging
# use the new "utils.shared_models" module for app model sharing
from utils.shared_models import (
update_model_symlinks, # main WORKER function (file/folder symlinks, Fix/remove broken symlinks, pull back local app models into shared)
SHARED_MODELS_DIR, SHARED_MODEL_FOLDERS, SHARED_MODEL_FOLDERS_FILE, ensure_shared_models_folders,
APP_INSTALL_DIRS, APP_INSTALL_DIRS_FILE, init_app_install_dirs, # APP_INSTALL_DIRS dict/file/function
MAP_APPS, sync_with_app_configs_install_dirs, # internal MAP_APPS dict and sync function
SHARED_MODEL_APP_MAP, SHARED_MODEL_APP_MAP_FILE, init_shared_model_app_map, # SHARED_MODEL_APP_MAP dict/file/function
write_dict_to_jsonfile, read_dict_from_jsonfile, PrettyDICT # JSON helper functions
)
# the "update_model_symlinks()" function replaces the app.py function with the same same
# and redirects to same function name "update_model_symlinks()" in the new "utils.shared_models" module
#
# this function does ALL the link management, including deleting "stale" symlinks,
# so the "recreate_symlinks()" function will be also re-routed to the
# "utils.shared_models.update_model_symlinks()" function (see CHANGE #3a and CHANGE #3b)
# the "ensure_shared_models_folders()" function will be called from app.py::create_shared_folders(),
# and replaces this function (see CHANGE #3)
# the "init_app_install_dirs() function initializes the
# global module 'APP_INSTALL_DIRS' dict: { 'app_name': 'app_installdir' }
# which does a default mapping from app code or (if exists) from external JSON 'APP_INSTALL_DIRS_FILE' file
# NOTE: this APP_INSTALL_DIRS dict is temporary synced with the 'app_configs' dict (see next)
# the "sync_with_app_configs_install_dirs() function syncs the 'APP_INSTALL_DIRS' dict's 'app_installdir' entries
# from the 'app_configs' dict's 'app_path' entries and uses the MAP_APPS dict for this task
# NOTE: this syncing is a temporary solution, and needs to be better integrated later
# the "init_shared_model_app_map()" function initializes the
# global module 'SHARED_MODEL_APP_MAP' dict: 'model_type' -> 'app_name:app_model_dir' (relative path)
# which does a default mapping from app code or (if exists) from external JSON 'SHARED_MODEL_APP_MAP_FILE' file
from utils.websocket_utils import send_websocket_message, active_websockets
from utils.app_configs import get_app_configs, add_app_config, remove_app_config, app_configs
from utils.model_utils import download_model, check_civitai_url, check_huggingface_url, SHARED_MODELS_DIR, format_size
from utils.model_utils import download_model, check_civitai_url, check_huggingface_url, format_size #, SHARED_MODELS_DIR # lutzapps - SHARED_MODELS_DIR is owned by shared_models module now
app = Flask(__name__)
sock = Sock(app)
@ -37,6 +73,8 @@ S3_BASE_URL = "https://better.s3.madiator.com/"
SETTINGS_FILE = '/workspace/.app_settings.json'
CIVITAI_TOKEN_FILE = '/workspace/.civitai_token'
HF_TOKEN_FILE = '/workspace/.hf_token' # lutzapps - added support for HF_TOKEN_FILE
def load_settings():
if os.path.exists(SETTINGS_FILE):
@ -91,6 +129,11 @@ def index():
pod_id=RUNPOD_POD_ID,
RUNPOD_PUBLIC_IP=os.environ.get('RUNPOD_PUBLIC_IP'),
RUNPOD_TCP_PORT_22=os.environ.get('RUNPOD_TCP_PORT_22'),
# lutzapps - CHANGE #2 - allow localhost Url for unsecure "http" and "ws" WebSockets protocol,
# according to LOCAL_DEBUG ENV var (used 3x in "index.html" changes)
enable_unsecure_localhost=os.environ.get('LOCAL_DEBUG'),
settings=settings,
current_auth_method=current_auth_method,
ssh_password=ssh_password,
@ -297,7 +340,20 @@ def remove_existing_app_config(app_name):
return jsonify({'status': 'success', 'message': f'App {app_name} removed successfully'})
return jsonify({'status': 'error', 'message': f'App {app_name} not found'})
# unused function
def obsolate_update_model_symlinks():
# lutzapps - CHANGE #3 - use the new "shared_models" module for app model sharing
# remove this whole now unused function
return "replaced by utils.shared_models.update_model_symlinks()"
# modified function
def setup_shared_models():
# lutzapps - CHANGE #4 - use the new "shared_models" module for app model sharing
jsonResult = update_model_symlinks()
return SHARED_MODELS_DIR # shared_models_dir is now owned and managed by the "shared_models" utils module
# remove below unused code
shared_models_dir = '/workspace/shared_models'
model_types = ['Stable-diffusion', 'VAE', 'Lora', 'ESRGAN']
@ -326,7 +382,12 @@ def setup_shared_models():
return shared_models_dir
def update_model_symlinks():
# unused function
def obsolate_update_model_symlinks():
# lutzapps - CHANGE #5 - use the new "shared_models" module for app model sharing
# remove this whole now unused function
return "replaced by utils.shared_models.update_model_symlinks()"
shared_models_dir = '/workspace/shared_models'
apps = {
'stable-diffusion-webui': '/workspace/stable-diffusion-webui/models',
@ -367,7 +428,7 @@ def update_model_symlinks():
print("Model symlinks updated.")
def update_symlinks_periodically():
while True:
while True:
update_model_symlinks()
time.sleep(300) # Check every 5 minutes
@ -375,7 +436,12 @@ def start_symlink_update_thread():
thread = threading.Thread(target=update_symlinks_periodically, daemon=True)
thread.start()
def recreate_symlinks():
# unused function
def obsolate_recreate_symlinks():
# lutzapps - CHANGE #6 - use the new "shared_models" module for app model sharing
# remove this whole now unused function
return "replaced by utils.shared_models.update_model_symlinks()"
shared_models_dir = '/workspace/shared_models'
apps = {
'stable-diffusion-webui': '/workspace/stable-diffusion-webui/models',
@ -421,16 +487,29 @@ def recreate_symlinks():
return "Symlinks recreated successfully."
# modified function
@app.route('/recreate_symlinks', methods=['POST'])
def recreate_symlinks_route():
# lutzapps - CHANGE #7 - use the new "shared_models" module for app model sharing
jsonResult = update_model_symlinks()
return jsonResult
# remove below unused code
try:
message = recreate_symlinks()
return jsonify({'status': 'success', 'message': message})
except Exception as e:
return jsonify({'status': 'error', 'message': str(e)})
# modified function
@app.route('/create_shared_folders', methods=['POST'])
def create_shared_folders():
# lutzapps - CHANGE #8 - use the new "shared_models" module for app model sharing
jsonResult = ensure_shared_models_folders()
return jsonResult
# remove below unused code
try:
shared_models_dir = '/workspace/shared_models'
model_types = ['Stable-diffusion', 'Lora', 'embeddings', 'VAE', 'hypernetworks', 'aesthetic_embeddings', 'controlnet', 'ESRGAN']
@ -467,11 +546,58 @@ def save_civitai_token(token):
with open(CIVITAI_TOKEN_FILE, 'w') as f:
json.dump({'token': token}, f)
# lutzapps - added function - 'HF_TOKEN' ENV var
def load_huggingface_token():
# look FIRST for Huggingface token passed in as 'HF_TOKEN' ENV var
HF_TOKEN = os.environ.get('HF_TOKEN', '')
if not HF_TOKEN == "":
print("'HF_TOKEN' ENV var found")
## send the found token to the WebUI "Models Downloader" 'hfToken' Password field to use
# send_websocket_message('extend_ui_helper', {
# 'cmd': 'hfToken', # 'hfToken' must match the DOM Id of the WebUI Password field in "index.html"
# 'message': "Put the HF_TOKEN in the WebUI Password field 'hfToken'"
# } )
return HF_TOKEN
# only if the 'HF_API_TOKEN' ENV var was not found, then handle it via local hidden HF_TOKEN_FILE
try:
if os.path.exists(HF_TOKEN_FILE):
with open(HF_TOKEN_FILE, 'r') as f:
data = json.load(f)
return data.get('token')
except:
return None
return None
# lutzapps - modified function - support 'CIVITAI_API_TOKEN' ENV var
def load_civitai_token():
if os.path.exists(CIVITAI_TOKEN_FILE):
with open(CIVITAI_TOKEN_FILE, 'r') as f:
data = json.load(f)
return data.get('token')
# look FIRST for CivitAI token passed in as 'CIVITAI_API_TOKEN' ENV var
CIVITAI_API_TOKEN = os.environ.get('CIVITAI_API_TOKEN', '')
if not CIVITAI_API_TOKEN == "":
print("'CIVITAI_API_TOKEN' ENV var found")
## send the found token to the WebUI "Models Downloader" 'hfToken' Password field to use
# send_websocket_message('extend_ui_helper', {
# 'cmd': 'civitaiToken', # 'civitaiToken' must match the DOM Id of the WebUI Password field in "index.html"
# 'message': 'Put the CIVITAI_API_TOKEN in the WebUI Password field "civitaiToken"'
# } )
return CIVITAI_API_TOKEN
# only if the 'CIVITAI_API_TOKEN' ENV var is not found, then handle it via local hidden CIVITAI_TOKEN_FILE
try:
if os.path.exists(CIVITAI_TOKEN_FILE):
with open(CIVITAI_TOKEN_FILE, 'r') as f:
data = json.load(f)
return data.get('token')
except:
return None
return None
@app.route('/save_civitai_token', methods=['POST'])
@ -487,18 +613,53 @@ def get_civitai_token_route():
token = load_civitai_token()
return jsonify({'token': token})
# lutzapps - add support for passed in "HF_TOKEN" ENV var
@app.route('/get_huggingface_token', methods=['GET'])
def get_hugginface_token_route():
token = load_huggingface_token()
return jsonify({'token': token})
# lutzapps - CHANGE #9 - return model_types to populate the Download manager Select Option
# new function to support the "Model Downloader" with the 'SHARED_MODEL_FOLDERS' dictionary
@app.route('/get_model_types', methods=['GET'])
def get_model_types_route():
model_types_dict = {}
# check if the SHARED_MODELS_DIR exists at the "/workspace" location!
# that only happens AFTER the the user clicked the "Create Shared Folders" button
# on the "Settings" Tab of the app's WebUI!
# to reload existing SHARED_MODEL_FOLDERS into the select options dropdown list,
# we send a WebSockets message to "index.html"
if not os.path.exists(SHARED_MODELS_DIR):
# return an empty model_types_dict, so the "Download Manager" does NOT get
# the already in-memory SHARED_MODEL_FOLDERS code-generated default dict
# BEFORE the workspace folders in SHARED_MODELS_DIR exists
return model_types_dict
i = 0
for model_type, model_type_description in SHARED_MODEL_FOLDERS.items():
model_types_dict[i] = {
'modelfolder': model_type,
'desc': model_type_description
}
i = i + 1
return model_types_dict
@app.route('/download_model', methods=['POST'])
def download_model_route():
url = request.json.get('url')
model_name = request.json.get('model_name')
model_type = request.json.get('model_type')
civitai_token = request.json.get('civitai_token') or load_civitai_token()
hf_token = request.json.get('hf_token')
hf_token = request.json.get('hf_token') or load_huggingface_token() # lutzapps - added HF_TOKEN ENV var support
version_id = request.json.get('version_id')
file_index = request.json.get('file_index')
is_civitai, _, _, _ = check_civitai_url(url)
is_huggingface, _, _, _, _ = check_huggingface_url(url)
is_huggingface, _, _, _, _ = check_huggingface_url(url) # TODO: double call
if not (is_civitai or is_huggingface):
return jsonify({'status': 'error', 'message': 'Unsupported URL. Please use Civitai or Hugging Face URLs.'}), 400
@ -507,7 +668,7 @@ def download_model_route():
return jsonify({'status': 'error', 'message': 'Civitai token is required for downloading from Civitai.'}), 400
try:
success, message = download_model(url, model_name, model_type, send_websocket_message, civitai_token, hf_token, version_id, file_index)
success, message = download_model(url, model_name, model_type, civitai_token, hf_token, version_id, file_index)
if success:
if isinstance(message, dict) and 'choice_required' in message:
return jsonify({'status': 'choice_required', 'data': message['choice_required']})
@ -522,7 +683,10 @@ def download_model_route():
@app.route('/get_model_folders')
def get_model_folders():
folders = {}
for folder in ['Stable-diffusion', 'VAE', 'Lora', 'ESRGAN']:
# lutzapps - replace the hard-coded model types
for folder, model_type_description in SHARED_MODEL_FOLDERS.items():
#for folder in ['Stable-diffusion', 'VAE', 'Lora', 'ESRGAN']:
folder_path = os.path.join(SHARED_MODELS_DIR, folder)
if os.path.exists(folder_path):
total_size = 0

0
official-templates/better-ai-launcher/gunicorn.conf.py → official-templates/better-ai-launcher/app/gunicorn.conf.py
View file

0
official-templates/better-ai-launcher/mushroom.png → official-templates/better-ai-launcher/app/static/mushroom.png
View file

Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 130 KiB

After

Width:  |  Height:  |  Size: 130 KiB

0
official-templates/better-ai-launcher/poddy-song.mp3 → official-templates/better-ai-launcher/app/static/poddy-song.mp3
View file

0
official-templates/better-ai-launcher/poddy.png → official-templates/better-ai-launcher/app/static/poddy.png
View file

Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 328 KiB

After

Width:  |  Height:  |  Size: 328 KiB

0
official-templates/better-ai-launcher/snake.png → official-templates/better-ai-launcher/app/static/snake.png
View file

Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 68 KiB

After

Width:  |  Height:  |  Size: 68 KiB

335
official-templates/better-ai-launcher/templates/index.html → official-templates/better-ai-launcher/app/templates/index.html
View file

@ -908,7 +908,8 @@
font-size: 12px;
}
#model-download-status,
/* lutzapps - allow word-wrap for model-download-status to see downloaded file at status: 'Complete'
/*#model-download-status,*/
#model-download-speed,
#model-download-eta {
margin-top: 5px;
@ -918,6 +919,12 @@
overflow: hidden;
text-overflow: ellipsis;
}
#model-download-status {
margin-top: 5px;
height: 20px;
line-height: 20px;
word-wrap: break-word;
}
/* Update the CSS for the token saving textbox */
#civitaiTokenSave {
@ -932,9 +939,9 @@
}
.example-urls {
background-color: #333;
border-radius: 5px;
padding: 15px;
background-color: #333;
border-radius: 5px;
padding: 15px;
}
.example-url {
@ -966,9 +973,9 @@
}
.model-folders-grid {
display: grid;
grid-template-columns: repeat(auto-fill, minmax(200px, 1fr));
gap: 15px;
display: grid;
grid-template-columns: repeat(auto-fill, minmax(200px, 1fr));
gap: 15px;
}
.model-folder {
@ -981,11 +988,12 @@
margin-top: 20px;
}
/* lutzapps - double definition
#model-download-status,
#model-download-speed,
#model-download-eta {
margin-top: 10px;
}
} */
#recreate-symlinks-container {
margin-top: 20px;
@ -1211,14 +1219,14 @@
<div class="model-downloader">
<input type="text" id="modelUrl" placeholder="Enter Civitai or Hugging Face URL">
<input type="text" id="modelName" placeholder="Enter model name (optional)">
<select id="modelType">
<option value="Stable-diffusion">Checkpoint (Stable Diffusion)</option>
<option value="Lora">LoRA</option>
<option value="VAE">VAE</option>
<option value="ESRGAN">Upscaler</option>
<select id="modelType"><!-- lutzapps - Change #1 - remove default modelType select option values -->
<!-- <option value="Stable-diffusion">Checkpoint (Stable Diffusion)</option> -->
<!-- <option value="Lora">LoRA</option> -->
<!-- <option value="VAE">VAE</option> -->
<!-- <option value="ESRGAN">Upscaler</option> -->
</select>
<input type="password" id="civitaiToken" placeholder="Civitai API Token">
<input type="text" id="hfToken" placeholder="Hugging Face API Token (optional)">
<input type="password" id="hfToken" placeholder="Hugging Face API Token (optional)">
<button onclick="downloadModel()" class="settings-button">Download Model</button>
</div>
<div id="model-download-progress" style="display: none;">
@ -1245,27 +1253,27 @@
<div class="example-urls">
<div class="example-url">
<span class="example-label">Stable Diffusion:</span>
<a href="#" class="example-link" id="example-sd" onclick="copyToClipboard(this.textContent); return false;"></a>
<a href="#" class="example-link" id="example-sd" onclick="useInModelDownloader('ckpt', this.textContent); return false;"></a>
</div>
<div class="example-url">
<span class="example-label">LoRA:</span>
<a href="#" class="example-link" id="example-lora" onclick="copyToClipboard(this.textContent); return false;"></a>
<a href="#" class="example-link" id="example-lora" onclick="useInModelDownloader('loras', this.textContent); return false;"></a>
</div>
<div class="example-url">
<span class="example-label">VAE:</span>
<a href="#" class="example-link" id="example-vae" onclick="copyToClipboard(this.textContent); return false;"></a>
<a href="#" class="example-link" id="example-vae" onclick="useInModelDownloader('vae', this.textContent); return false;"></a>
</div>
<div class="example-url">
<span class="example-label">Upscaler:</span>
<a href="#" class="example-link" id="example-upscaler" onclick="copyToClipboard(this.textContent); return false;"></a>
<a href="#" class="example-link" id="example-upscaler" onclick="useInModelDownloader('upscale_models', this.textContent); return false;"></a>
</div>
<div class="example-url">
<span class="example-label">Flux Dev:</span>
<a href="#" class="example-link" id="example-flux-dev" onclick="copyToClipboard(this.textContent); return false;"></a>
<a href="#" class="example-link" id="example-flux-dev" onclick="useInModelDownloader('unet', this.textContent); return false;"></a>
</div>
<div class="example-url">
<span class="example-label">Flux Schnell:</span>
<a href="#" class="example-link" id="example-flux-schnell" onclick="copyToClipboard(this.textContent); return false;"></a>
<a href="#" class="example-link" id="example-flux-schnell" onclick="useInModelDownloader('unet', this.textContent); return false;"></a>
</div>
</div>
</div>
@ -1302,7 +1310,34 @@
let currentLogAppName = null;
const podId = '{{ pod_id }}';
const WS_PORT = 7222; // This is the Nginx port
const WS_URL = `wss://${podId}-${WS_PORT}.proxy.runpod.net/ws`;
// lutzapps - remember the last know MODELTYPE_SELECTED of the WebUI Dom Id 'modelType' "select dropdown" model list
// so the selection can be restored between "Tab Switches", and during refreshes/reloads of the modelType list
// this is handled by the extendUIHelper() function, which is called directly from JavaScript events,
// but is also called indirectly from Python code via WS message type='extend_ui_helper'
// e.g. from the model_utils:download_civitai_model() function, to preserve "state" of the selected modelType
let MODELTYPE_SELECTED = "";
modelType.onchange = function() {
MODELTYPE_SELECTED = this.value;
alert("onchange=" + MODELTYPE_SELECTED);
}
// *** lutzapps - Change #2 - support to run locally at http://localhost:${WS_PORT} (3 locations in "index.html")
const enable_unsecure_localhost = '{{ enable_unsecure_localhost }}';
// default is to use the "production" WeckSockets CloudFlare URL
// NOTE: ` (back-ticks) are used here for template literals
var WS_URL = `wss://${podId}-${WS_PORT}.proxy.runpod.net/ws`; // need to be declared as var
if (`${enable_unsecure_localhost}` === 'True') { // value of LOCAL_DEBUG ENV var
// make sure to use "ws" Protocol (insecure) instead of "wss" (WebSockets Secure) for localhost,
// otherwise you will get the 'loadingOverlay' stay and stays on screen with ERROR:
// "WebSocket disconnected. Attempting to reconnect..." blocking the webpage http://localhost:7222
WS_URL = `ws://localhost:${WS_PORT}/ws`; // localhost WS (unsecured)
//alert(`Running locally with WS_URL=${WS_URL}`);
}
let socket;
let reconnectAttempts = 0;
@ -1407,6 +1442,9 @@
}
if (data.stage === 'Complete') {
// lutzapps - clear 'Speed: N/A' and 'ETA: Calculating...' Div textContent, when no eta/speed data, but 'Complete'
etaDiv.textContent = '';
speedDiv.textContent = '';
loadModelFolders(); // Refresh the Existing Models section when download is complete
}
}
@ -1446,7 +1484,14 @@
}
function openApp(appKey, port) {
const url = `https://${podId}-${port}.proxy.runpod.net/`;
// *** lutzapps - Change #3 - support to run locally
// NOTE: ` (back-ticks) are used here for template literals
var url = `https://${podId}-${port}.proxy.runpod.net/`; // need to be declared as var
if (`${enable_unsecure_localhost}` === 'True') {
url = `http://localhost:${port}/`; // remove runpod.net proxy
//alert(`openApp URL=${url}`);
}
window.open(url, '_blank');
}
@ -1510,34 +1555,170 @@
updateStatus();
setInterval(updateStatus, 5000);
setInterval(updateLogs, 1000);
var data = {};
data.cmd = 'refreshModelTypes';
extendUIHelper(data); // lutzapps - initialize the available SHARED_MODEL_FOLDERS for the "Model Downloader" modelType select list
}
// lutzapps - populate modeltype select options from shared_models
async function extendUIHelper(data) {
// check the data
// if no data is passed, the default cmd is "selectModelType" with no specific "model_type",
// which means to re-select the last selected modelType option from the global var MODELTYPE_SELECTED
var cmd = "selectModelType"; // default cmd, when called with empty data or empty cmd
var model_type = MODELTYPE_SELECTED; // the 'modelType' option value (=foldername) for the select dropdown list
var token = ""; // the token value of HF_TOKEN or CIVITAI_API_TOKEN to pass to the corresponding WebUI Password fields on the "Models" tab
var response;
var result;
if (data !== undefined && data.cmd !== undefined) {
cmd = data.cmd;
if (cmd === "selectModelType" && data.model_type !== undefined) {
model_type = data.model_type; // the model_type which is passed in to select
}
if ((cmd === "hfToken" || cmd === "civitaiToken") // cmd need to match the DOM id of the Password field on the "Models" tab
&& data.token !== undefined) {
token = data.token; // if token = undefined or empty "", then the corresponding token get fetched from the server
}
}
//alert("extendUIHelper(): cmd=" + cmd +", model_type=" + model_type + ", token=" + token); // debug-info (DISABLED)
switch (cmd) {
case "civitaiToken":
if (token === "") { // get the data from the Server
response = await fetch('/get_civitai_token');
result = await response.json();
token = result['token'];
}
alert(cmd + "=" + token);
// pass tokens from HF or CIVITAI ENV vars into their Password fields
document.getElementById(cmd).value = token; //'********'; // indicate a found token, but require to call back
break;
case "hfToken":
if (token === "") { // get the data from the Server
response = await fetch('/get_huggingface_token');
result = await response.json();
token = result['token'];
}
alert(cmd + "=" + token);
// pass tokens from HF or CIVITAI ENV vars into their Password fields
document.getElementById(cmd).value = token; //'********'; // indicate a found token, but require to call back
break;
case "refreshModelTypes":
// refresh and optionally select the 'modelType' list for "Model Downloader"
var modelTypeSelect = document.getElementById('modelType');
// get the data from the Server
response = await fetch('/get_model_types');
result = await response.json();
//alert(JSON.stringify(result)); // show the JSON-String
var model_types = result; // get the JSON-Object
var count = Object.keys(model_types).length; // count=18, when using the default SHARED_MODEL_FOLDERS dict
// the "/get_model_types" app.get_model_types_route() function checks
// if the SHARED_MODELS_DIR shared files already exists at the "/workspace" location.
// that only happens AFTER the the user clicked the "Create Shared Folders" button
// on the "Settings" Tab of the app's WebUI.
// it will return an empty model_types_dict, so the "Download Manager" does NOT get
// the already in-memory SHARED_MODEL_FOLDERS code-generated default dict
// BEFORE the workspace folders in SHARED_MODELS_DIR exists!
//
// when SHARED_MODELS_DIR exists (or updates), this function will be called via a Socket Message
// to "refresh" its content automatically
var modelTypeSelected = modelTypeSelect.value; // remember the current selected modelType.option value
modelTypeSelect.options.length = 0; // clear all current modelTypeSelect options
for (i = 0 ; i < count; i += 1) {
modelTypeOption = document.createElement('option');
modelType = model_types[String(i)]['modelfolder'];
modelTypeOption.setAttribute('value', modelType);
modelTypeOption.appendChild(document.createTextNode(model_types[String(i)]['desc']));
//if (modelFolder === modelTypeSelected) {
// modelTypeOption.selected = true; // reselect it
//}
modelTypeSelect.appendChild(modelTypeOption);
}
//modelTypeSelect.selectedIndex = modelfolder_index; // set the selected index
//modelTypeSelect.options[mmodelfolder_index].selected = true; // and mark it as "selected" option
if (modelTypeSelected === "") { // initial refresh, called by initializeUI() function
modelTypeSelect.selectedIndex = 0; // use the first modelType option, usually "ckpt"
MODELTYPE_SELECTED = modelTypeSelect.options[0].value; // NOT handled by the onchange() event handler
}
else {
modelTypeSelect.value = modelTypeSelected; // (re-)apply the selected modelType option
MODELTYPE_SELECTED = modelTypeSelected; // NOT handled by the onchange() event handler
}
break;
case "selectModelType":
// this is called by model_utils:download_civitai_model() which passed the downloading 'model_type'
// to select for clarity, which can be different from the currently selected modelType option
// if called without a 'model_type', the last MODELTYPE_SELECTED will be selected (after a "Tab Switch")
// refresh and optionally select the 'modelType' list for "Model Downloader"
var modelTypeSelect = document.getElementById('modelType');
modelTypeSelect.value = model_type;
MODELTYPE_SELECTED = model_type; // NOT handled by the onchange() event handler
break;
default: // no cmd passed is same as "selectModelType" without a 'model_type'
// this is already handled by "selectModelType" defaults, there is no "default" case needed here
}
}
async function downloadModel() {
const url = document.getElementById('modelUrl').value;
const modelName = document.getElementById('modelName').value;
const modelType = document.getElementById('modelType').value;
const hfToken = document.getElementById('hfToken').value;
let civitaiToken = null;
// Check if the URL is from Civitai
if (url.includes('civitai.com')) {
civitaiToken = document.getElementById('civitaiToken').value;
if (civitaiToken === '********') {
// If the token is masked, fetch it from the server
const response = await fetch('/get_civitai_token');
const result = await response.json();
civitaiToken = result.token;
const url = document.getElementById('modelUrl').value;
const modelName = document.getElementById('modelName').value;
const modelType = document.getElementById('modelType').value;
//const hfToken = document.getElementById('hfToken').value;
let hfToken = null;
let civitaiToken = null;
// lutzapps - support HF_TOKEN ENV var
// Check if the URL is from Huggingface
if (url.toLowerCase().includes('huggingface.co')) { // be case-insensitive with this url
hfToken = document.getElementById('hfToken').value;
if (hfToken === '********') {
// if the token is masked, fetch it from the server
const response = await fetch('/get_huggingface_token');
const result = await response.json();
hfToken = result.token;
}
}
// Check if the URL is from Civitai
if (url.toLowerCase().includes('civitai.com')) { // lutzapps - be case-insensitive with this url
civitaiToken = document.getElementById('civitaiToken').value;
if (civitaiToken === '********') {
// If the token is masked, fetch it from the server
const response = await fetch('/get_civitai_token');
const result = await response.json();
civitaiToken = result.token;
}
if (!civitaiToken) {
alert('Civitai API token is required for Civitai URLs. Please enter it in the Civitai API Token field.');
return;
}
}
if (!civitaiToken) {
alert('Civitai API token is required for Civitai URLs. Please enter it in the Civitai API Token field.');
return;
}
await startModelDownload(url, modelName, modelType, civitaiToken, hfToken);
}
await startModelDownload(url, modelName, modelType, civitaiToken, hfToken);
}
async function startModelDownload(url, modelName, modelType, civitaiToken, hfToken, versionId = null, fileIndex = null) {
@ -1739,6 +1920,19 @@
}
}
// lutzapps - added HF_TOKEN support
async function loadHFToken() {
try {
const response = await fetch('/get_huggingface_token');
const result = await response.json();
if (result.token) {
document.getElementById('hfToken').value = '********';
}
} catch (error) {
console.error('Error loading Huggingface token:', error);
}
}
function formatSize(sizeInBytes) {
const units = ['B', 'KB', 'MB', 'GB', 'TB'];
let size = sizeInBytes;
@ -1765,7 +1959,7 @@
}
}
// Example URLs
// Example URLs - lutzapps - TODO: enrich data
const exampleUrls = {
'Stable-diffusion': 'https://civitai.com/models/90352/dreamshaper',
'Lora': 'https://civitai.com/models/58390?modelVersionId=62833',
@ -1784,6 +1978,25 @@
document.getElementById('example-flux-schnell').textContent = exampleUrls['Flux-Schnell'];
}
// lutzapps - replace function copyToClipboard() with function useInModelDownloader()
function useInModelDownloader(modelType, modelUrl) {
// copy the downloadUrl in the Model Downloader "Url" textbox
document.getElementById('modelUrl').value = modelUrl;
// select the modelType in the modelType select list, we are just about to download
var data = {};
data.cmd = 'selectModelType';
data.model_type = modelType;
extendUIHelper(data);
navigator.clipboard.writeText(modelUrl).then(() => {
alert('URL copied to Downloader and into the Clipboard!');
}, (err) => {
console.error('Could not copy text: ', err);
});
}
// lutzapps - obsolete function (can be deleted)
function copyToClipboard(text) {
navigator.clipboard.writeText(text).then(() => {
alert('URL copied to clipboard!');
@ -1794,8 +2007,12 @@
// Call this function when the Models tab is opened
document.querySelector('.navbar-tabs a[onclick="openTab(event, \'models-tab\')"]').addEventListener('click', function() {
loadModelFolders();
//alert("querySelector");
loadModelFolders(); // lutzapps - this ModelFolders is NOT for the 'modelType' "select dropdown" model list
extendUIHelper(); // lutzapps - select the last know MODELTYPE_SELECTED in the WebUI Dom Id 'modelType' "select dropdown" model list
loadCivitaiToken();
loadHFToken(); // lutzapps - added HF_TOKEN ENV var Support
updateExampleUrls();
});
@ -1820,10 +2037,13 @@
// Additional actions when switching tabs
if (tabName === 'apps-tab') {
document.querySelector('.logs-section').style.display = 'flex';
updateLogs(); // Refresh logs when switching back to Apps tab
updateLogs(); // Refresh logs when switching back to the "Apps" tab
} else if (tabName === 'models-tab') {
loadModelFolders();
loadCivitaiToken();
// lutzapps - the following event handler functions already fired
// in the queryselector eventhandler (esspecially for the "models-tab" tab)
//loadModelFolders();
//loadCivitaiToken();
//loadHFToken(); // lutzapps - added HF_TOKEN Support
} else if (tabName === 'settings-tab') {
loadSshDetails();
updateFileBrowserStatus();
@ -1914,11 +2134,19 @@
}
function openFileBrowser() {
const podId = '{{ pod_id }}';
const url = `https://${podId}-7222.proxy.runpod.net/fileapp/`;
window.open(url, '_blank');
const podId = '{{ pod_id }}';
// *** lutzapps - Change #5 - support to run locally
// NOTE: ` (back-ticks) are used here for template literals
var url = `https://${podId}-7222.proxy.runpod.net/fileapp/`; // need to be declared as var
if (`${enable_unsecure_localhost}` === 'True') {
url = `http://localhost:7222/fileapp/`; // remove runpod.net proxy
//alert(`FileBrowser URL=${url}`);
}
window.open(url, '_blank');
}
async function controlFileBrowser(action) {
try {
const response = await fetch(`/${action}_filebrowser`);
@ -2015,6 +2243,9 @@
updateModelDownloadProgress(data.data);
} else if (data.type === 'status_update') {
updateAppStatus(data.data);
// lutzapps - use the extendUIHelper to "bridge" certain extensions between Python and JavaScript
} else if (data.type === 'extend_ui_helper') {
extendUIHelper(data.data);
}
// Handle other message types as needed
} catch (error) {

469
official-templates/better-ai-launcher/app/tests/README-SHARED_MODELS.txt Normal file
View file

@ -0,0 +1,469 @@
TESTDATA AND EXPLANATION OF MAPPING EVERYTHING YOU WANT
In the folder "/app/tests" you find the following files:
/app/tests/
- "README-SHARED_MODELS.txt" (this file)
- "populate_testdata.sh" (bash script to un-tar and expand all testdata into the "/workspace" folder)
- "testdata_shared_models_link.tar.gz" (Testcase #1, read below)
- "testdata_stable-diffusion-webui_pull.tar.gz" (Testcase #2, read below)
- "testdata_installed_apps_pull.tar.gz" (Testcase #3, read below)
CREATE TESTDATA (once done already):
cd /workspace
# For Testcase #1 - create testdata in "shared_models" folder with dummy models for most model_types:
$ tar -czf testdata_shared_models_link.tar.gz shared_models
# For Testcase #2 - create testdata with SD-Models for A1111 to be pulled back into "shared_models" and linked back:
$ tar -czf testdata_stable-diffusion-webui_pull.tar.gz stable-diffusion-webui
# For Testcase #3 -create testdata with all possible "Apps" installed into your "/workspace"
$ tar -czf /app/tests/testdata_installed_apps_pull.tar.gz Apps
USE TESTDATA:
# BEFORE(!) you run this script, read the readme ;-)
/app/tests/populate_testdata.sh::
# use these 3 test cases and extract/merge them accordingly into your workspace, bur READ before you mess you up too much!!
tar -xzf /app/tests/testdata_shared_models_link.tar.gz /workspace
tar -xzf /app/tests/testdata_stable-diffusion-webui_pull.tar.gz /workspace
tar -xzf /app/tests/testdata_installed_apps_pull.tar.gz /workspace
Testcase #1:
When you expand "./testdata_shared_models_link.tar.gz" into the "/workspace" folder, you get:
$ tree shared_models
shared_models
├── LLM
│   └── Meta-Llama-3.1-8B
│   ├── llm-Llama-modelfile1.txt
│   ├── llm-Llama-modelfile2.txt
│   └── llm-Llama-modelfile3.txt
├── ckpt
│   ├── ckpt-model1.txt
│   └── ckpt-model2.txt
├── clip
│   └── clip-model1.txt
├── controlnet
│   └── controlnet-model1.txt
├── embeddings
│   ├── embedding-model1.txt
│   └── embedding-model2.txt
├── hypernetworks
│   └── hypernetworks-model1.txt
├── insightface
│   └── insightface-model1.txt
├── ipadapters
│   ├── ipadapter-model1.txt
│   └── xlabs
│   └── xlabs-ipadapter-model1.txt
├── loras
│   ├── flux
│   │   └── flux-lora-model1.txt
│   ├── lora-SD-model1.txt
│   ├── lora-SD-model2.txt
│   ├── lora-SD-model3.txt
│   ├── lora-SD-model4.txt
│   ├── lora-SD-model5.txt
│   ├── lora-model1.txt
│   ├── lora-model2.txt
│   └── xlabs
│   └── xlabs-lora-model1.txt
├── reactor
│   ├── faces
│   │   └── reactor-faces-model1.txt
│   └── reactor-model1.txt
├── unet
│   ├── unet-model1.txt
│   └── unet-model2.txt
├── upscale_models
│   └── esrgan-model1.txt
├── vae
│   └── vae-model1.txt
└── vae-approx
└── vae-apporox-model1.txt
20 directories, 29 files
All these "*.txt" files "simulate" model files of a specific category (model type).
When you have this test data and you click the "Recreate Symlinks" button on the "Settings" Tab, all these models will be shared with all "installed" apps, like:
A1111: /workspace/stable-diffusion-webui
Forge: /workspace/stable-diffusion-webui-forge
ComfyUI: /workspace/ComfyUI
Kohya_ss: /workspace/Kohya_ss
CUSTOM1: /workspace/joy-caption-batch
To "simulate" the installed app, you just need to create one or all of these folders manually, as empty folders. Maybe try it one-by-one, like you would do "in-real-life".
After there is at least ONE app installed, you can test the model sharing with the above mentioned Button "Recreate Symlinks".
All of these 29 models should be shared into all "installed" apps.
When you "add" a second app, also this new app will get all these models shared into its model folders, which can be differently named.
Some model types (e.g. UNET) have a separate model folder from Checkpoints in ComfyUI, but in A1111/Forge, these 2 model types will be merged ("flattened") in one "Stable-Diffusion" model folder. See later in the third MAP shown here.
The same goes in the other direction.
LoRA models which are "organized" in subfolders like "flux" or "slabs" in shared models folder, or in ComfyUI will again be flattened for app like A1111/Forge into the only ONE "Lora" folder.
Pay attention to the details, as the Folder is called "Lora" (no "s" and capital "L" for A111/Forge), but is called "Loras" (with ending "s" and all lower-case).
You can even map/share LLM models, which mostly come installed as a folder with many files needed for one LLM. For these models, FOLDER symlinks will be created instead or regular file symlinks which are sufficient for most model files of many model types.
Some model types, like "Embeddings" are managed differently for A1111/Forge (outside its regular model folder), which is not the case for ComfyUI. All this can be tested.
You can also test to delete a "shared model" in the "shared_models directory, and all its "symlinked" copies should be also automatically removed.
Testcase #2:
In the second testdata TAR archive, you have some SD-models which simulate the installation of the model files once for ONE app, in this test case only for A1111.
The "./testdata_stable-diffusion-webui_pull.tar.gz" is easier to handle than the one for Testcase #3, as it installs directly into the "original" App install location.
$ tree stable-diffusion-webui
stable-diffusion-webui
├── _add
│   ├── lora-SD-model2.txt
│   ├── lora-SD-model3.txt
│   ├── lora-SD-model4.txt
│   └── lora-SD-model5.txt
└── models
└── Lora
└── lora-SD-model1.txt
4 directories, 5 files
Testcase #3:
In this test case you also have other apps installed already, but the principle is the same, just a little bit more careful folder management.
The "./testdata_installed_apps_pull.tar.gz.tar.gz" extracts into an "Apps" folder.
All folders in this extracted "Apps" folder should be copied into the "/workspace" folder, to simulate an installed A1111, Forge, ComfyUI and Kohya_ss. Make sure that at the end you NOT see the extracted "Apps" folder anymore, as you only used its SUB-FOLDERS to copy/move them into "/workspace" and replace/merge existing folder.
$ tree Apps
Apps
├── ComfyUI
├── Kohya_ss
├── _add
│   ├── lora-SD-model2.txt
│   ├── lora-SD-model3.txt
│   ├── lora-SD-model4.txt
│   └── lora-SD-model5.txt
├── joy-caption-batch
├── stable-diffusion-webui
│   └── models
│   └── Lora
│   └── lora-SD-model1.txt
└── stable-diffusion-webui-forge
9 directories, 5 files
This test cases #2 and #3 start with one SD LoRA model "lora-SD-model1.txt" installed only for "stable-diffusion-webui" (A1111) in its "Lora" model folder.
Wenn you have this test data installed, and you click the "Recreate Symlinks" button of the "better-ai-laucncher" template, then this "locally" (one-app-only) installed LoRA model will be found and "pulled" back into the "shared_models/loras" folder, and "re-shared" back to the pulled location in A1111.
But it will then also be shared to all other "installed" apps, like ComfyUI, Forge. But not into Kohya, as there is no "mapping rule" defined for Kohya for LoRA models.
The only "mapping rule" for Kohya which is defined, is to get all "ckpt" (Checkpoint) model files and all UNET model files shared from the corresponding "shared_models" subfolders into its /models folder (see later in the 3rd MAP below).
In the testdata "Apps" folder you also find an "_add" folder, with 4 more SD-Models to play around with the App Sharing/Syncing framework. Put the in any local app model folder and watch what happens to them and where they the can be seen/used from other apps. You either wait a fewMinutes to let this happen automatically (every 5 Minutes), or you press the "Recreate Symlinks" button at any time to kick this off.
You can also test to see what happens, when you DELETE a model file from the shared_models sub-folders, and that all its symlinks shared to all apps will also automatically be removed, so no broken links will be left behind.
When you delete a symlink in an app model folder, only the local app "looses" the model (it is just only a link to the original shared model), so no worries here. Such locally removed symlinks however will be re-created again automatically.
BUT THAT IS ONLY THE BEGINNING.
All this logic described here is controlled via 3 (three) "MAP" dictionary JSON files, which can be found also in the "/workspace/shared_models" folder, after you click the "Create Shared Folders" button on the "Settings" Tab. They will be auto-generated, if missing, or otherwise "used-as-is" :
1.) The "SHARED_MODEL_FOLDERS" map found as "SHARED_MODEL_FOLDERS_FILE"
"/workspace/shared_models/_shared_model_folders.json":
{
# "model_type" (=subdir_name of SHARED_MODELS_DIR): "model_type_description"
"ckpt": "Model Checkpoint (Full model including a CLIP and VAE model)",
"clip": "CLIP Model (used together with UNET models)",
"controlnet": "ControlNet model (Canny, Depth, Hed, OpenPose, Union-Pro, etc.)",
"embeddings": "Embedding (aka Textual Inversion) Model",
"hypernetworks": "HyperNetwork Model",
"insightface": "InsightFace Model",
"ipadapters": "ControlNet IP-Adapter Model",
"ipadapters/xlabs": "IP-Adapter from XLabs-AI",
"LLM": "LLM (aka Large-Language Model) is folder mapped (1 folder per model), append '/*' in the map",
"loras": "LoRA (aka Low-Ranking Adaption) Model",
"loras/xlabs": "LoRA Model from XLabs-AI",
"loras/flux": "LoRA Model trained on Flux.1 Dev or Flux.1 Schnell",
"reactor": "Reactor Model",
"reactor/faces": "Reactor Face Model",
"unet": "UNET Model Checkpoint (need separate CLIP and VAE Models)",
"upscale_models": "Upscaling Model (based on ESRGAN)",
"vae": "VAE En-/Decoder Model",
"vae-approx": "Approximate VAE Model"
}
This is the "SHARED_MODEL_FOLDERS" map for the "shared_models" sub-folders, which are used by the App Model Sharing, and this is also the dictionary which is used by the "Model Downloader" Dropdown Listbox, which can be found on the "Models" Tab of the WebUI.
The idea is to make it easy for people to download their models from "Huggingface" or "CivitAI" directly into the right model type subfolder of the "/workspace/shared_models" main folder.
This "SHARED_MODEL_FOLDERS" map also is used when you click the "Create Shared Folders" button.
NOTE: pay special attention to the examples with the "loras" folder. There is one regular model type for "loras", and 2 "grouping" model types, "loras/flux" and "loras/xlabs".
We come back to "grouping" mapping rules later, when we discuss the 3rd map.
Feel free to add/remove/edit/rename items here, alls should be re-created automatically. Just nothing from renamed folders will be deleted.
NOTE: if you change something in this map file, you need to "read" it into the App via the "Create Shared Folders" button on the "Settings" Tab of the WebUI.
If new folders are found in the map, they will be created, but nothing you have already created before will be deleted or renamed automatically, so be careful not generating two folders for the same model type, or move the models manually into the renamed folder.
IMPORTANT: Also be aware that when you add or change/rename folder names here, you need to also add or change/rename these folder names in the third "SHARED_MODEL_APP_MAP" explained below!!!
2.) The "APP_INSTALL_DIRS" map found as "APP_INSTALL_DIRS_FILE"
"/workspace/shared_models/_app_install_dirs.json":
{
# "app_name": "app_install_dir"
"A1111": "/workspace/stable-diffusion-webui",
"Forge": "/workspace/stable-diffusion-webui-forge",
"ComfyUI": "/workspace/ComfyUI",
"Kohya_ss": "/workspace/Kohya_ss",
"CUSTOM1": "/workspace/joy-caption-batch"
}
This is the "APP_INSTALL_DIRS" map for the app install dirs within the "/workspace", and as you see, it also supports "CUSTOM" apps to be installed and participating at the model sharing.
This dictionary is "synced" with the main apps "app_configs" dictionary, so the installation folders are the same, and this should NOT be changed. What you can change in this MAP is to add "CUSTOM" apps, like "CUSTOM1" here e.g. to re-use the Llama LLM model which is centrally installed in "shared_models" under the LLM folder to be "shared" between ComfyUI and "Joy Caption Batch" tool, which is nice to generate your "Caption" files for your LoRA Training files with "Kohya_ss" for example.
3.) The "SHARED_MODEL_APP_MAP" map found as "SHARED_MODEL_APP_MAP_FILE"
"/workspace/shared_models/_shared_model_app_map.json":
{
"ckpt": { # "model_type" (=subdir_name of SHARED_MODELS_DIR)
# "app_name": "app_model_folderpath" (for this "model_type", path is RELATIVE to "app_install_dir" of APP_INSTALL_DIRS map)
"ComfyUI": "/models/checkpoints",
"A1111": "/models/Stable-diffusion",
"Forge": "/models/Stable-diffusion",
"Kohya_ss": "/models" # flatten all "ckpt" / "unet" models here
},
"clip": {
"ComfyUI": "/models/clip",
"A1111": "/models/text_encoder",
"Forge": "/models/text_encoder"
},
"controlnet": {
"ComfyUI": "/models/controlnet",
"A1111": "/models/ControlNet",
"Forge": "/models/ControlNet"
#"A1111": "/extensions/sd-webui-controlnet/models", # SD1.5 ControlNets
#"Forge": "/extensions/sd-webui-controlnet/models" # SD1.5 ControlNets
},
# EMBEDDINGS map outside of models folder for FORGE / A1111
"embeddings": {
"ComfyUI": "/models/embeddings",
"A1111": "/embeddings",
"Forge": "/embeddings"
},
"hypernetworks": {
"ComfyUI": "/models/hypernetworks",
"A1111": "/models/hypernetworks",
"Forge": "/models/hypernetworks"
},
"insightface": {
"ComfyUI": "/models/insightface",
"A1111": "/models/insightface", # unverified location
"Forge": "/models/insightface" # unverified location
},
"ipadapters": {
"ComfyUI": "/models/ipadapter/",
"A1111": "/extensions/sd-webui-controlnet/models", # unverified location
"Forge": "/extensions/sd-webui-controlnet/models" # unverified location
},
"ipadapters/xlabs": { # sub-folders for XLabs-AI IP-Adapters
"ComfyUI": "/models/xlabs/ipadapters",
"A1111": "/extensions/sd-webui-controlnet/models", # flatten all "xlabs" ipadapters here
"Forge": "/extensions/sd-webui-controlnet/models" # flatten all "xlabs" ipadapters here
},
# some LoRAs get stored here in sub-folders, e.g. "/xlabs/*"
"loras": {
"ComfyUI": "/models/loras",
"A1111": "/models/Lora",
"Forge": "/models/Lora"
},
# Support "XLabs-AI" LoRA models
"loras/xlabs": { # special syntax for "grouping"
"ComfyUI": "/models/loras/xlabs",
"A1111": "/models/Lora", # flatten all "xlabs" LoRAs here
"Forge": "/models/Lora" # flatten all "xlabs" LoRAs here
},
# Support "Grouping" all FLUX LoRA models into a LoRA "flux" sub-folder for ComfyUI,
# which again need to be flattened for other apps
"loras/flux": {
"ComfyUI": "/models/loras/flux",
"A1111": "/models/Lora", # flatten all "flux" LoRAs here
"Forge": "/models/Lora" # flatten all "flux" LoRAs here
},
"reactor": {
"ComfyUI": "/models/reactor", # unverified location
"A1111": "/models/reactor",
"Forge": "/models/reactor",
},
"reactor/faces": {
"ComfyUI": "/models/reactor/faces", # unverified location
"A1111": "/models/reactor",
"Forge": "/models/reactor",
},
# UNET models map into the CKPT folders of all other apps, except for ComfyUI
"unet": {
"ComfyUI": "/models/unet",
"A1111": "/models/Stable-diffusion", # flatten all "ckpts" / "unet" models here
"Forge": "/models/Stable-diffusion", # flatten all "ckpts" / "unet" models here
"Kohya_ss": "/models" # flatten all "ckpt" / "unet" models here
},
"upscale_models": {
"ComfyUI": "/models/upscale_models",
"A1111": "/models/ESRGAN",
"Forge": "/models/ESRGAN"
},
"vae": {
"ComfyUI": "/models/vae",
"A1111": "/models/VAE",
"Forge": "/models/VAE"
},
"vae-approx": {
"ComfyUI": "/models/vae_approx",
"A1111": "/models/VAE-approx",
"Forge": "/models/VAE-approx"
},
# E.g. Custom Apps support for Joytag-Caption-Batch Tool (which uses the "Meta-Llama-3.1-8B" LLM)
# to share the model with e.g. ComfyUI. This LLM model come as full folders with more than one file!
# Pay attention to the special syntax for folder mappings (add a "/*" suffix to denote a folder mapping)
"LLM/Meta-Llama-3.1-8B/*": { # special syntax for "folder" symlink (the "/*" is mandatory)
"ComfyUI": "/models/LLM/Meta-Llama-3.1-8B/*", # special syntax for "folder" symlink, the "/*" is optional
"CUSTOM1": "/model/*" # special syntax for "folder" symlink, the "/*" is optional
}
}
This third and last "SHARED_MODEL_APP_MAP" "connects" the "SHARED_MODEL_FOLDERS" map with the "APP_INSTALL_DIRS" map, to have a very flexible "Mapping" between all apps.
You can custom define all folder/directory layouts and namings of model types, which you can think of.
Add new mappings to your likings.
And if you don't need some of these mappings, then change them, or delete them.
This is sample data to give you a head start what you can do.
NOTE: as already introduced before, with the 3 "loras" model type folders, here we now see the 2 grouping model types "/loras/flux" and "/loras/xlabs" now applied in a "grouping" map rule.
"Grouping" of specific model.
E.g. LoRAs into separate sub-folders for LoRAs (e.g. "flux", "xlabs") for easier "filtering" in ComfyUI. Look this MAP and the testdata for LoRA samples of that feature.
It shows it for 2 LoRA sub-folders as just mentioned. These 2 "grouping" map rules also show how to "flatten" these sub-folders into only one "Lora" model folder for A1111/Forge, all of them go flat into their "Lora" folder, as these are apps, that do NOT support sub-folders per model type, and need therfore to have these "flux" and "xlabs" LoRAs "flattended" to see and consume them.
Try to add your own "grouping" map rule, or delete "grouping" map rules you not need.
Otherwise this "SHARED_MODEL_APP_MAP" should be self-explanatory, except for the last part with the FOLDER sharing syntax, as shown here and already mentioned LLM "Meta-Llama-3.1-8B", which will be used here, to show how an LLM model as "Meta-Llama-3.1-8B" can be shared between the app "ComfyUI" and a "CUSTOM1" defined app "joy-caption-batch".
FOLDER SHARING, e.g. LLM folder-based models:
_app_install.dirs.json:
{
...
"Kohya_ss": "/workspace/Kohya_ss",
"CUSTOM1": "/workspace/joy-caption-batch"
}
_shared_model_folders.json:
{
...
"LLM": "LLM (aka Large-Language Model) is folder mapped (1 folder per model), append '/*' in the map",
}
To define a "folder" map rule, the "rule" must be an EXISTING shared foldername trailing with "/*",
e.g.
_shared_model_app_map.json:
{
...
"LLM/Meta-Llama-3.1-8B/*": {
"ComfyUI": "/models/LLM/Meta-Llama-3.1-8B/*",
"CUSTOM1": "/model/*"
}
}
This difference in the syntax with a trailing "/*" shows the shared_models Framework, that you NOT want the model files IN the folder to be shared one-by-one (which would not work anyway for a LLM model), but that you want to share the whole FOLDER, so EVERYTHING within the folder is also automatically shared.
Just append a "/*" to the MAP rule of the physical folder name "LLM/Meta-Llama-3.1-8B", and it will understand, that you want a FOLDER symlink from the folder path with the "/*" removed.
All pathes are relative to "/workspace/shared_models".
This will "trigger" a "folder" symlink to all defined target app folder mapped folders.
The target foldernames (defined with or without the trailing "/*") must be a NON-EXISTING app foldername, which will be the target folder symlink from the shared folder source.
Here are the detailes of this folder symlink for "ComfyUI":
$ cd ComfyUI/models/LLM
$ ls -la
drwxr-xr-x 3 root root 96 Oct 25 11:45 .
drwxr-xr-x 18 root root 576 Oct 25 11:45 ..
lrwxr-xr-x 1 root root 46 Oct 25 11:45 Meta-Llama-3.1-8B -> /workspace/shared_models/LLM/Meta-Llama-3.1-8B
NOTE: pay special attention to the second folder map rule for the "CUSTOM1" app "joy-cation-batch". The target mapped folder "model" has a different name from the source folder "LLM/Meta-Llama-3.1-8B/", and the 3 sample LLM model files all go directly into the linked "/model" folder.
Here are the detailes of this folder symlink for "CUSTOM1" (joy-caption-batch):
$ cd joy-caption-batch
$ ls -la
drwxr-xr-x 3 root root 96 Oct 25 11:53 .
drwxr-xr-x 10 root root 320 Oct 25 11:45 ..
lrwxr-xr-x 1 root root 46 Oct 25 11:53 model -> /workspace/shared_models/LLM/Meta-Llama-3.1-8B
$ cd model
$ ls -la
drwxr-xr-x 6 root root 192 Oct 23 10:22 .
drwxr-xr-x 5 root root 160 Oct 23 10:27 ..
-rwx------ 1 root root 51 Oct 18 11:29 llm-Llama-modelfile1.txt
-rwx------ 1 root root 51 Oct 18 11:29 llm-Llama-modelfile2.txt
-rwx------ 1 root root 51 Oct 18 11:29 llm-Llama-modelfile3.txt
Folder Sharing rules is an advanced technique, and they can only be manually added, editing the "SHARED_MODEL_APP_MAP" file for such rules, as shown above.
While all new downloaded "single" model files of a model type will be automatically shared into all app model folders, "folder models" (as typically found for LLM models), need to be added to this "SHARED_MODELS_MAP" JSON file manually to be shared, as it is shown in this example.
You could also use this folder mappings for "custom" LoRA sub-categories in ComfyUI "per folder", instead of sub-folders with separate file symlinks, but I not want to go deeper here.
You can try it, nothing bad will happen.
If you delete one of these three MAP JSON files, they will be re-generated with its shown "default" content.
But when you edit/change these 3 MAP JSON files, they will be used with your changes/addings every time you use this "/workspace" volume.
That's it for the FULL version of "shared_models" sharing between apps, and how you can easily test this before using the "real data" ;-)
Regards,
lutzapps

12
official-templates/better-ai-launcher/app/tests/populate_testdata.sh Executable file
View file

@ -0,0 +1,12 @@
#!/bin/bash
# please use the "./readme-testdata.txt" before you extract these TARs!!!
# Testcase #1
tar -xzf /app/tests/testdata_shared_models_link.tar.gz /workspace
# Testcase #2
tar -xzf /app/tests/testdata_stable-diffusion-webui_pull.tar.gz /workspace
# Testcase #3
tar -xzf /app/tests/testdata_installed_apps_pull.tar.gz /workspace

BIN
official-templates/better-ai-launcher/app/tests/testdata_installed_apps_pull.tar.gz Normal file
View file

Binary file not shown.

BIN
official-templates/better-ai-launcher/app/tests/testdata_shared_models_link.tar.gz Normal file
View file

Binary file not shown.

BIN
official-templates/better-ai-launcher/app/tests/testdata_stable-diffusion-webui_pull.tar.gz Normal file
View file

Binary file not shown.

0
official-templates/better-ai-launcher/utils/app_configs.py → official-templates/better-ai-launcher/app/utils/app_configs.py
View file

7
official-templates/better-ai-launcher/utils/app_utils.py → official-templates/better-ai-launcher/app/utils/app_utils.py
View file

@ -300,7 +300,12 @@ def install_app(app_name, app_configs, send_websocket_message):
else:
return False, f"Unknown app: {app_name}"
def update_model_symlinks():
# unused function
def onsolate_update_model_symlinks():
# lutzapps - CHANGE #7 - use the new "shared_models" module for app model sharing
# remove this whole now unused function
return "replaced by utils.shared_models.update_model_symlinks()"
shared_models_dir = '/workspace/shared_models'
apps = {
'stable-diffusion-webui': '/workspace/stable-diffusion-webui/models',

0
official-templates/better-ai-launcher/utils/filebrowser_utils.py → official-templates/better-ai-launcher/app/utils/filebrowser_utils.py
View file

488
official-templates/better-ai-launcher/app/utils/model_utils.py Normal file
View file

@ -0,0 +1,488 @@
import os
import requests
from urllib.parse import unquote, urlparse
from tqdm import tqdm
import json
import re
import time
import math
### model_utils-v0.2 by lutzapps, Oct 30th 2024 ###
# lutzapps - modify for new shared_models module and overwrite for this module
from utils.shared_models import (ensure_shared_models_folders, update_model_symlinks, SHARED_MODELS_DIR)
from utils.websocket_utils import send_websocket_message, active_websockets
#SHARED_MODELS_DIR = '/workspace/shared_models' # this global var is now owned by the 'shared_models' module
# lutzapps - modify this CivitAI model_type mapping to the new SHARED_MODEL_FOLDERS map
MODEL_TYPE_MAPPING = {
# CivitAI-Modeltype: SHARED_MODEL_FOLDERS
'Checkpoint': 'ckpt', #'Stable-diffusion', # not clear name for model_type
'LORA': 'loras', #'Lora', # now lowercase and plural
'LoCon': 'loras', #'Lora', # now lowercase and plural
'TextualInversion': 'embeddings',
'VAE': 'vae', #'VAE', # now lowercase
'Hypernetwork': 'hypernetworks',
'AestheticGradient': 'embeddings', #'aesthetic_embeddings', # store together with "embeddings"
'ControlNet': 'controlnet',
'Upscaler': 'upscale_models' #'ESRGAN' # there are probably other upscalers not based on ESRGAN
}
def ensure_shared_folder_exists():
# lutzapps - replace with new shared_models code
#for folder in ['Stable-diffusion', 'Lora', 'embeddings', 'VAE', 'hypernetworks', 'aesthetic_embeddings', 'controlnet', 'ESRGAN']:
# os.makedirs(os.path.join(SHARED_MODELS_DIR, folder), exist_ok=True)
ensure_shared_models_folders()
def check_civitai_url(url):
prefix = "civitai.com"
try:
if prefix in url:
if "civitai.com/api/download" in url:
version_id = url.strip("/").split("/")[-1]
return False, True, None, int(version_id)
subpath = url[url.find(prefix) + len(prefix):].strip("/")
url_parts = subpath.split("?")
if len(url_parts) > 1:
model_id = url_parts[0].split("/")[1]
version_id = url_parts[1].split("=")[1]
return True, False, int(model_id), int(version_id)
else:
model_id = subpath.split("/")[1]
return True, False, int(model_id), None
except (ValueError, IndexError):
print("Error parsing Civitai model URL")
return False, False, None, None
def check_huggingface_url(url):
parsed_url = urlparse(url)
if parsed_url.netloc not in ["huggingface.co", "huggingface.com"]:
return False, None, None, None, None
path_parts = [p for p in parsed_url.path.split("/") if p]
if len(path_parts) < 5 or (path_parts[2] != "resolve" and path_parts[2] != "blob"):
return False, None, None, None, None
repo_id = f"{path_parts[0]}/{path_parts[1]}"
branch_name = path_parts[3]
remaining_path = "/".join(path_parts[4:])
folder_name = os.path.dirname(remaining_path) if "/" in remaining_path else None
filename = unquote(os.path.basename(remaining_path))
return True, repo_id, filename, folder_name, branch_name
def download_model(url, model_name, model_type, civitai_token=None, hf_token=None, version_id=None, file_index=None):
ensure_shared_folder_exists()
is_civitai, is_civitai_api, model_id, _ = check_civitai_url(url)
is_huggingface, repo_id, hf_filename, hf_folder_name, hf_branch_name = check_huggingface_url(url) # TODO: double call
if is_civitai or is_civitai_api:
if not civitai_token:
return False, "Civitai token is required for downloading from Civitai"
success, message = download_civitai_model(url, model_name, model_type, civitai_token, version_id, file_index)
elif is_huggingface:
success, message = download_huggingface_model(url, model_name, model_type, repo_id, hf_filename, hf_folder_name, hf_branch_name, hf_token)
else:
return False, "Unsupported URL"
if success:
send_websocket_message('model_download_progress', {
'percentage': 100,
'stage': 'Complete',
'message': 'Download complete and symlinks updated'
})
return success, message
# lutzapps - added SHA256 checks for already existing ident and downloaded HuggingFace model
def download_civitai_model(url, model_name, model_type, civitai_token, version_id=None, file_index=None):
try:
is_civitai, is_civitai_api, model_id, url_version_id = check_civitai_url(url)
headers = {'Authorization': f'Bearer {civitai_token}'}
if is_civitai_api:
api_url = f"https://civitai.com/api/v1/model-versions/{url_version_id}"
else:
api_url = f"https://civitai.com/api/v1/models/{model_id}"
response = requests.get(api_url, headers=headers)
response.raise_for_status()
model_data = response.json()
if is_civitai_api:
version_data = model_data
model_data = version_data['model']
else:
if version_id:
version_data = next((v for v in model_data['modelVersions'] if v['id'] == version_id), None)
elif url_version_id:
version_data = next((v for v in model_data['modelVersions'] if v['id'] == url_version_id), None)
else:
version_data = model_data['modelVersions'][0]
if not version_data:
return False, f"Version ID {version_id or url_version_id} not found for this model."
civitai_model_type = model_data['type']
model_type = MODEL_TYPE_MAPPING.get(civitai_model_type, 'Stable-diffusion')
files = version_data['files']
if file_index is not None and 0 <= file_index < len(files):
file_to_download = files[file_index]
elif len(files) > 1:
# If there are multiple files and no specific file was chosen, ask the user to choose
file_options = [{'name': f['name'], 'size': f['sizeKB'], 'type': f['type']} for f in files]
return True, {
'choice_required': {
'type': 'file',
'model_id': model_id,
'version_id': version_data['id'],
'files': file_options
}
}
else:
civitai_file = files[0] # that is the metadata civitai_file
download_url = civitai_file['downloadUrl']
if not model_name:
model_name = civitai_file['name']
model_path = os.path.join(SHARED_MODELS_DIR, model_type, model_name)
platformInfo = {
"platform_name": 'civitai',
"civitai_file": civitai_file # civitai_file metadata dictionary
}
# call shared function for "huggingface" and "civitai" for SHA256 support and "Model Downloader UI" extended support
download_sha256_hash, found_ident_local_model, message = get_modelfile_hash_and_ident_existing_modelfile_exists(
model_name, model_type, model_path, # pass local workspace vars, then platform specific vars as dictionary
platformInfo) # [str, bool, str]
if found_ident_local_model:
return True, message
# model_path does NOT exist - run with original code
os.makedirs(os.path.dirname(model_path), exist_ok=True)
# lutzapps - add SHA256 check for download_sha256_hash is handled after download finished in download_file()
return download_file(download_url, download_sha256_hash, model_path, headers) # [bool, str]
except Exception as e: # requests.RequestException as e:
return False, f"Exception downloading from CivitAI: {str(e)}"
# lutzapps - calculate the SHA256 hash string of a file
def get_sha256_hash_from_file(file_path:str) -> tuple[bool, str]:
import hashlib # support SHA256 checks
try:
sha256_hash = hashlib.sha256()
with open(file_path, "rb") as f:
# read and update hash string value in blocks of 4K
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
return True, sha256_hash.hexdigest().upper()
except Exception as e:
return False, str(e)
# lutzapps - support SHA256 Hash check of already locally existing modelfile against its metadata hash before downloading is needed
# shared function for "huggingface" and "civitai" called by download_huggingface_model() and download_civitai_model()
def get_modelfile_hash_and_ident_existing_modelfile_exists(model_name:str, model_type:str, model_path:str, platformInfo:dict) -> tuple[bool, str, str]:
try:
# update (and remember) the selected index of the modelType select list of the "Model Downloader"
message = f"Select the ModelType '{model_type}' to download"
print(message)
send_websocket_message('extend_ui_helper', {
'cmd': 'selectModelType',
'model_type': f'{model_type}', # e.g. "loras" or "vae"
'message': message
} )
# get the SHA256 hash - used for compare against existing or downloaded model
platform_name = platformInfo['platform_name'].lower() # currently "civitai" or "huggingface", but could be extendend
print(f"\nPlatform: {platform_name}")
match platform_name:
case "huggingface":
# get the platform-specific passed variables for "huggingface"
hf_token = platformInfo['hf_token']
repo_id = platformInfo['repo_id']
hf_filename = platformInfo['hf_filename']
#from huggingface_hub import hf_hub_download
# lutzapps - to get SHA256 hash from model
from huggingface_hub import (
# HfApi, # optional when not calling globally
get_paths_info #list_files_info #DEPRECATED/MISSING: list_files_info => get_paths_info
)
from huggingface_hub.hf_api import (
RepoFile, RepoFolder, BlobLfsInfo
)
## optionally configure a HfApi client instead of calling globally
# hf_api = HfApi(
# endpoint = "https://huggingface.co", # can be a Private Hub endpoint
# token = hf_token, # token is not persisted on the machine
# )
print(f"getting SHA256 Hash for '{model_name}' from repo {repo_id}/{hf_filename}")
# HfApi.list_files_info deprecated -> HfApi.get_paths_info (runs into exception, as connot be imported as missing)
#files_info = hf_api.list_files_info(repo_id, hf_filename, expand=True)
#paths_info = hf_api.get_paths_info(repo_id, hf_filename, expand=True) # use via HfApi
paths_info = get_paths_info(repo_id, hf_filename, expand=True) # use global (works fine)
repo_file = paths_info[0] # RepoFile or RepoFolder class instance
# check for RepoFolder or NON-LFS
if isinstance(repo_file, RepoFolder):
raise NotImplementedError("Downloading a folder is not implemented.")
if not repo_file.lfs:
raise NotImplementedError("Copying a non-LFS file is not implemented.")
lfs = repo_file.lfs # BlobLfsInfo class instance
download_sha256_hash = lfs.sha256.upper()
print(f"Metadata from RepoFile LFS '{repo_file.rfilename}'")
print(f"SHA256: {download_sha256_hash}")
case "civitai":
# get the platform-specific passed variables for "civitai"
civitai_file = platformInfo['civitai_file'] # civitai_file metadata dictionary
# get the SHA256 hash - used for compare against existing or downloaded model
download_sha256_hash = civitai_file['hashes']['SHA256'] # civitai_file = passed file
### END platform specific code
# check if model file already exists
if not os.path.exists(model_path):
message = f"No local model '{os.path.basename(model_path)}' installed"
print(message)
return download_sha256_hash, False, message
message = f"Model already exists: {os.path.basename(model_path)}, SHA256 check..."
print(message)
send_websocket_message('model_download_progress', {
'percentage': 0, # ugly
'stage': 'Downloading',
'message': message
})
# check if existing model is ident with model to download
# this can *take a while* for big models, but even better than to unnecessarily redownload the model
successfull_HashGeneration, model_sha256_hash = get_sha256_hash_from_file(model_path)
# if NOT successful, the hash contains the Exception
print(f"SHA256 hash generated from local file: '{model_path}'\n{model_sha256_hash}")
if successfull_HashGeneration and model_sha256_hash == download_sha256_hash:
message = f"Existing and ident model aleady found for '{os.path.basename(model_path)}'"
print(message)
send_websocket_message('model_download_progress', {
'percentage': 100,
'stage': 'Complete',
'message': message
})
return download_sha256_hash, successfull_HashGeneration, message
else:
if successfull_HashGeneration: # the generated SHA256 file model Hash did not match against the metadata hash
message = f"Local installed model '{os.path.basename(model_path)}' has DIFFERENT \nSHA256: {model_sha256_hash}"
print(message)
return download_sha256_hash, False, message
else: # NOT successful, the hash contains the Exception
error_msg = model_sha256_hash
error_msg = f"Exception occured while generating the SHA256 hash for '{model_path}':\n{error_msg}"
print(error_msg)
except Exception as e:
error_msg = f"Exception when downloading from {platform_name}: {str(e)}"
return "", False, error_msg # hash, identfile, message
# lutzapps - added SHA256 checks for already existing ident and downloaded HuggingFace model
def download_huggingface_model(url, model_name, model_type, repo_id, hf_filename, hf_folder_name, hf_branch_name, hf_token=None):
try:
from huggingface_hub import hf_hub_download
if not model_name:
model_name = hf_filename
model_path = os.path.join(SHARED_MODELS_DIR, model_type, model_name)
platformInfo = {
"platform_name": 'huggingface',
"hf_token": hf_token,
"repo_id": repo_id,
"hf_filename": hf_filename
}
# call shared function for "huggingface" and "civitai" for SHA256 support and "Model Downloader UI" extended support
download_sha256_hash, found_ident_local_model, message = get_modelfile_hash_and_ident_existing_modelfile_exists(
model_name, model_type, model_path, # pass local workspace vars, then platform specific vars as dictionary
platformInfo) # [str, bool, str]
if found_ident_local_model:
return True, message
# model_path does NOT exist - run with original code
os.makedirs(os.path.dirname(model_path), exist_ok=True)
send_websocket_message('model_download_progress', {
'percentage': 0,
'stage': 'Downloading',
'message': f'Starting download from Hugging Face: {repo_id}'
})
kwargs = {
'repo_id': repo_id,
'filename': hf_filename,
'subfolder': hf_folder_name,
'revision': hf_branch_name,
'local_dir': os.path.dirname(model_path)
#'local_dir_use_symlinks': False # deprecated, should be removed
}
if hf_token:
kwargs['token'] = hf_token
file_path = hf_hub_download(**kwargs) ### HF_DOWNLOAD_START
### HF_DOWNLOAD COMPLETE
# SHA256 Hash checks of downloaded modelfile against its metadata hash
# call shared function for "huggingface" and "civitai" for SHA256 support and "Model Downloader UI" extended support
return check_downloaded_modelfile(file_path, download_sha256_hash, "huggingface") # [bool, str]
except Exception as e:
return False, f"Exception when downloading from 'HuggingFace': {str(e)}"
# lutzapps - added SHA256 check for downloaded CivitAI model
def download_file(url, download_sha256_hash, file_path, headers=None):
try:
response = requests.get(url, stream=True, headers=headers)
response.raise_for_status()
total_size = int(response.headers.get('content-length', 0))
block_size = 8192
downloaded_size = 0
start_time = time.time()
with open(file_path, 'wb') as file: ### CIVITAI_DOWNLOAD
for data in response.iter_content(block_size):
size = file.write(data)
downloaded_size += size
current_time = time.time()
elapsed_time = current_time - start_time
if elapsed_time > 0:
speed = downloaded_size / elapsed_time
percentage = (downloaded_size / total_size) * 100 if total_size > 0 else 0
eta = (total_size - downloaded_size) / speed if speed > 0 else 0
send_websocket_message('model_download_progress', {
'percentage': round(percentage, 2),
'speed': f"{speed / (1024 * 1024):.2f} MB/s",
'eta': int(eta),
'stage': 'Downloading',
'message': f'Downloaded {format_size(downloaded_size)} / {format_size(total_size)}'
})
### CIVITAI_DOWNLOAD COMPLETE
# SHA256 Hash checks of downloaded modelfile against its metadata hash
# call shared function for "huggingface" and "civitai" for SHA256 support and "Model Downloader UI" extended support
return check_downloaded_modelfile(file_path, download_sha256_hash, "civitai") # [bool, str]
except Exception as e:
return False, f"Exception when downloading from CivitAI: {str(e)}"
# lutzapps - SHA256 Hash checks of downloaded modelfile against its metadata hash
# shared function for "huggingface" and "civitai" for SHA256 support and "Model Downloader UI" extended support
def check_downloaded_modelfile(model_path:str, download_sha256_hash:str, platform_name:str) -> tuple[bool, str]:
try:
# lutzapps - SHA256 check for download_sha256_hash
if download_sha256_hash == "":
return False, f"Downloaded model could not be verified with Metadata, no SHA256 hash found on '{platform_name}'"
# check if downloaded local model file is ident with HF model download_sha256_hash metadata
# this can take a while for big models, but even better than to have a corrupted model
send_websocket_message('model_download_progress', {
'percentage': 90, # change back from 100 to 90 (ugly)
'stage': 'Complete', # leave it as 'Complete' as this "clears" SPEED/ETA Divs
'message': f'SHA256 Check for Model: {os.path.basename(model_path)}'
})
successfull_HashGeneration, model_sha256_hash = get_sha256_hash_from_file(model_path)
if successfull_HashGeneration and model_sha256_hash == download_sha256_hash:
send_websocket_message('model_download_progress', {
'percentage': 100,
'stage': 'Complete',
'message': f'Download complete: {os.path.basename(model_path)}'
})
update_model_symlinks() # create symlinks for this new downloaded model for all installed apps
return True, f"Successfully downloaded (SHA256 checked, and symlinked) '{os.path.basename(model_path)}' from {platform_name}"
else:
if successfull_HashGeneration: # the generated SHA256 file model Hash did not match against the metadata hash
message = f"The downloaded model '{os.path.basename(model_path)}' has DIFFERENT \nSHA256: {model_sha256_hash} as stored on {platform_name}\nFile is possibly corrupted and was DELETED!"
print(message)
os.remove(model_path) # delete corrupted, downloaded file
return download_sha256_hash, False, message
else: # NOT successful, the hash contains the Exception
error_msg = model_sha256_hash
error_msg = f"Exception occured while generating the SHA256 hash for '{model_path}':\n{error_msg}"
print(error_msg)
except Exception as e:
error_msg = f"Exception when downloading from {platform_name}: {str(e)}"
return False, error_msg
# smaller helper functions
def get_civitai_file_size(url, token):
headers = {'Authorization': f'Bearer {token}'}
try:
response = requests.head(url, headers=headers, allow_redirects=True)
return int(response.headers.get('content-length', 0))
except:
return None
def get_huggingface_file_size(repo_id, filename, folder_name, branch_name, token):
from huggingface_hub import hf_hub_url, HfApi
try:
api = HfApi()
file_info = api.hf_hub_url(repo_id, filename, subfolder=folder_name, revision=branch_name)
response = requests.head(file_info, headers={'Authorization': f'Bearer {token}'} if token else None)
return int(response.headers.get('content-length', 0))
except:
return None
def format_size(size_in_bytes):
if size_in_bytes == 0:
return "0 B"
size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
i = int(math.floor(math.log(size_in_bytes, 1024)))
p = math.pow(1024, i)
s = round(size_in_bytes / p, 2)
return f"{s} {size_name[i]}"

921
official-templates/better-ai-launcher/app/utils/shared_models.py Normal file
View file

@ -0,0 +1,921 @@
import os
import shutil
import datetime
import threading
import time
import json
from flask import jsonify
from utils.websocket_utils import send_websocket_message, active_websockets
from utils.app_configs import (get_app_configs)
### shared_models-v0.9.1 by lutzapps, Oct 30th 2024 ###
### dev-my-v0.6
# to run (and optionally DEBUG) this docker image "better-ai-launcher" in a local container on your own machine
# you need to define the ENV var "LOCAL_DEBUG" in the "VSCode Docker Extension"
# file ".vscode/tasks.json" in the ENV settings of the "dockerRun" section (or any other way),
# and pass into the docker container:
# tasks.json:
# ...
# "dockerRun": {
# "containerName": "madiator2011-better-launcher", // no "/" allowed here for container name
# "image": "madiator2011/better-launcher:dev",
# "envFiles": ["${workspaceFolder}/.env"], // pass additional env-vars (hf_token, civitai token, ssh public-key) from ".env" file to container
# "env": { // this ENV vars go into the docker container to support local debugging
# "LOCAL_DEBUG": "True", // change app to localhost Urls and local Websockets (unsecured)
# "FLASK_APP": "app/app.py",
# "FLASK_ENV": "development", // changed from "production"
# "GEVENT_SUPPORT": "True" // gevent monkey-patching is being used, enable gevent support in the debugger
# // "FLASK_DEBUG": "0" // "1" allows debugging in Chrome, but then VSCode debugger not works
# },
# "volumes": [
# {
# "containerPath": "/app",
# "localPath": "${workspaceFolder}" // the "/app" folder (and sub-folders) will be mapped locally for debugging and hot-reload
# },
# {
# "containerPath": "/workspace",
# // TODO: create this folder before you run!
# "localPath": "${userHome}/Projects/Docker/Madiator/workspace"
# }
# ],
# "ports": [
# {
# "containerPort": 7222, // main Flask app port "AppManager"
# "hostPort": 7222
# },
# ...
#
#
# NOTE: to use the "LOCAL_DEBUG" ENV var just for local consumption of this image, you can run it like
#
# docker run -it -d --name madiator-better-launcher -p 22:22 -p 7777:7777 -p 7222:7222 -p 3000:3000 -p 7862:7862 -p 7863:7863
# -e LOCAL_DEBUG="True" -e RUNPOD_PUBLIC_IP="127.0.0.1" -e RUNPOD_TCP_PORT_22="22"
# -e PUBLIC_KEY="ssh-ed25519 XXXXXXX...XXXXX user@machine-DNS.local"
# --mount type=bind,source=/Users/test/Projects/Docker/madiator/workspace,target=/workspace
# madiator2011/better-launcher:dev
#
# To run the full 'app.py' / 'index.html' webserver locally, is was needed to "patch" the
# '/app/app.py' main application file and the
# '/app/templates/index.html'file according to the "LOCAL_DEBUG" ENV var,
# to switch the CF "proxy.runpod.net" Url for DEBUG:
#
### '/app/app.py' CHANGES ###:
# # lutzapps - CHANGE #1
# LOCAL_DEBUG = os.environ.get('LOCAL_DEBUG', 'False') # support local browsing for development/debugging
# ...
# filebrowser_status = get_filebrowser_status()
# return render_template('index.html',
# apps=app_configs,
# app_status=app_status,
# pod_id=RUNPOD_POD_ID,
# RUNPOD_PUBLIC_IP=os.environ.get('RUNPOD_PUBLIC_IP'),
# RUNPOD_TCP_PORT_22=os.environ.get('RUNPOD_TCP_PORT_22'),
# # lutzapps - CHANGE #2 - allow localhost Url for unsecure "http" and "ws" WebSockets protocol,
# according to LOCAL_DEBUG ENV var (used 3x in "index.html" changes)
# enable_unsecure_localhost=os.environ.get('LOCAL_DEBUG'),
# ...
# other (non-related) app.py changes omitted here
#
### '/app/template/index.html' CHANGES ###:
# <script>
# ...
# // *** lutzapps - Change #2 - support to run locally at http://localhost:${WS_PORT} (3 locations in "index.html")
# const enable_unsecure_localhost = '{{ enable_unsecure_localhost }}';
#
# // default is to use the "production" WeckSockets CloudFlare URL
# // NOTE: ` (back-ticks) are used here for template literals
# var WS_URL = `wss://${podId}-${WS_PORT}.proxy.runpod.net/ws`; // need to be declared as var
#
# if `${enable_unsecure_localhost}` === 'True') { // value of LOCAL_DEBUG ENV var
# // make sure to use "ws" Protocol (insecure) instead of "wss" (WebSockets Secure) for localhost,
# // otherwise you will get the 'loadingOverlay' stay and stays on screen with ERROR:
# // "WebSocket disconnected. Attempting to reconnect..." blocking the webpage http://localhost:7222
# WS_URL = `ws://localhost:${WS_PORT}/ws`; // localhost WS (unsecured)
# //alert(`Running locally with WS_URL=${WS_URL}`);
# }
# ...
# function openApp(appKey, port) {
# // *** lutzapps - Change #3 - support to run locally
# // NOTE: ` (back-ticks) are used here for template literals
# var url = `https://${podId}-${port}.proxy.runpod.net/`; // need to be declared as var
# if `${enable_unsecure_localhost}` === 'True') {
# url = `http://localhost:${port}/`; // remove runpod.net proxy
# //alert(`openApp URL=${url}`);
# }
#
# window.open(url, '_blank');
# }
# ...
# function openFileBrowser() {
# const podId = '{{ pod_id }}';
#
# // *** lutzapps - Change #5 - support to run locally
# // NOTE: ` (back-ticks) are used here for template literals
# var url = `https://${podId}-7222.proxy.runpod.net/fileapp/`; // need to be declared as var
# if `${enable_unsecure_localhost}` === 'True') {
# url = `http://localhost:7222/fileapp/`; // remove runpod.net proxy
# //alert(`FileBrowser URL=${url}`);
# }
#
# window.open(url, '_blank');
# }
# other (non-related) index.html changes omitted here
README_FILE_PREFIX = "_readme-" # prefix for all different dynamically generated README file names
### support local docker container runs with locally BOUND Workspace, needed also during local debugging
SHARED_MODELS_DIR = "/workspace/shared_models" # storage root for all shared_models (as designed for production app)
# check for "DISABLE_PULLBACK_MODELS" ENV var, and convert it from String to Boolean
# DISABLE_PULLBACK_MODELS: if not present (or empty) then "pull-back" local model files is enabled [default]
# else local found model files in app model folders will NOT be pulled back (and re-shared) into 'shared_models'
DISABLE_PULLBACK_MODELS = False
disable_pullback = os.environ.get('DISABLE_PULLBACK_MODELS', 'False').lower() # "True" / "true" / "1", or "False" / "false" / "0" / "" (not set)
if not (disable_pullback == "" or disable_pullback == "false" or disable_pullback == "0"):
DISABLE_PULLBACK_MODELS = True
# check for "DISABLE_PULLBACK_MODELS" ENV var, and convert it from String to Boolean
# LOCAL_DEBUG: if not present (or empty) then run in "production" [default],
# else run as "development" / "debug" version
LOCAL_DEBUG = False
local_debug_str = os.environ.get('LOCAL_DEBUG', 'False').lower() # "True" / "true" / "1", or "False" / "false" / "0" / "" (not set)
if not (local_debug_str == "" or local_debug_str == "false" or local_debug_str == "0"):
LOCAL_DEBUG = True
# show current configuration
print("\n\n*** SHARED_MODELS module init ***\n")
print(f"SHARED_MODELS_DIR='{SHARED_MODELS_DIR}'\n")
print(f"LOCAL_DEBUG='{LOCAL_DEBUG}'")
print(f"DISABLE_PULLBACK_MODELS='{DISABLE_PULLBACK_MODELS}'")
# show/hide the 3 dictionary MAPPING FILES, which control the whole module
MAKE_MAPPING_FILES_HIDDEN = False # can also be set according to LOCAL_DEBUG = True or False
if MAKE_MAPPING_FILES_HIDDEN:
HIDDEN_FILE_PREFIX = "." # filenames starting with a dot (".") are hidden by the filesystem
else:
HIDDEN_FILE_PREFIX = "" # filenames are shown to the user
print(f"MAKE_MAPPING_FILES_HIDDEN='{MAKE_MAPPING_FILES_HIDDEN}'\n")
# helper function to return a pretty formatted DICT string for human consumption (Logs, JSON)
def PrettyDICT(dict:dict) -> str:
dict_string = json.dumps(dict, ensure_ascii=False, indent=4, separators=(',', ': '))
return dict_string
# helper function called by init_shared_model_app_map() and init_shared_models_folders()
def write_dict_to_jsonfile(dict:dict, json_filepath:str, overwrite:bool=False) -> bool:
# Convert the 'dict' to JSON, and write the JSON object to file 'json_filepath'
#json_string = json.dumps(dict, indent=4, ensure_ascii=False, sort_keys=True)
try:
if os.path.exists(json_filepath) and not overwrite:
error_msg = f"dictionary file '{json_filepath}' already exists (and overwrite={overwrite})"
#print(error_msg)
return False, error_msg # failure
# Write the JSON data to a file
with open(json_filepath, 'w', encoding='utf-8') as output_file:
json.dump(dict, output_file, ensure_ascii=False, indent=4, separators=(',', ': '))
except Exception as e:
error_msg = f"ERROR in shared_models:write_dict_to_jsonfile() - loading JSON Map File '{json_filepath}'\nException: {str(e)}"
print(error_msg)
return False, error_msg # failure
return True, "" # success
# helper function called by init_shared_model_app_map() and init_shared_models_folders()
def read_dict_from_jsonfile(json_filepath:str) -> dict:
# Read JSON file from 'json_filepath' and return it as 'dict'
try:
if os.path.exists(json_filepath):
with open(json_filepath, 'r') as input_file:
dict = json.load(input_file)
else:
error_msg = f"dictionary file '{json_filepath}' does not exist"
#print(error_msg)
return {}, error_msg # failure
except Exception as e:
error_msg = f"ERROR in shared_models:read_dict_from_jsonfile() - loading JSON Map File '{json_filepath}'\nException: {str(e)}"
print(error_msg)
return {}, error_msg # failure
return dict, "" # success
# helper function for "init_app_install_dirs(), "init_shared_model_app_map()" and "init_shared_models_folders()"
def init_global_dict_from_file(dict:dict, dict_filepath:str, dict_description:str) -> bool:
# load or initialize the 'dict' for 'dict_description' from 'dict_filepath'
try:
if not os.path.exists(SHARED_MODELS_DIR):
print(f"\nThe SHARED_MODELS_DIR '{SHARED_MODELS_DIR}' is not found!\nCreate it by clicking the 'Create Shared Folders' button from the WebUI 'Settings' Tab\n")
return
if os.path.isfile(dict_filepath) and os.path.exists(dict_filepath):
dict_filepath_found = True
# read the dict_description from JSON file
print(f"\nExisting '{dict_description}' found and read from file '{dict_filepath}'\nThe file overwrites the code defaults!")
dict, error_msg = read_dict_from_jsonfile(dict_filepath)
if not error_msg == "":
print(error_msg)
else: # init the dict_description from app code
dict_filepath_found = False
print(f"No '{dict_description}'_FILE found, initializing default '{dict_description}' from code ...")
# use already defined dict from app code
# write the dict to JSON file
success, ErrorMsg = write_dict_to_jsonfile(dict, dict_filepath)
if success:
print(f"'{dict_description}' is initialized and written to file '{dict_filepath}'")
else:
print(ErrorMsg)
# Convert 'dict_description' dictionary to formatted JSON
print(f"\nUsing {'external' if dict_filepath_found else 'default'} '{dict_description}':\n{PrettyDICT(dict)}")
except Exception as e:
error_msg = f"ERROR in shared_models:init_global_dict_from_file() - initializing dict Map File '{dict_filepath}'\nException: {str(e)}"
print(error_msg)
return False, error_msg
return True, "" # success
# the below SHARED_MODEL_FOLDERS_FILE will be read and used (if exists),
# otherwise this file will be generated with the content of the below default SHARED_MODEL_FOLDERS dict
SHARED_MODEL_FOLDERS_FILE = f"{SHARED_MODELS_DIR}/{HIDDEN_FILE_PREFIX}_shared_model_folders.json"
SHARED_MODEL_FOLDERS = {
# "model_type" (=subdir_name of SHARED_MODELS_DIR): "model_type_description"
"ckpt": "Model Checkpoint (Full model including a CLIP and VAE model)",
"clip": "CLIP Model (used together with UNET models)",
"controlnet": "ControlNet model (Canny, Depth, Hed, OpenPose, Union-Pro, etc.)",
"embeddings": "Embedding (aka Textual Inversion) Model",
"hypernetworks": "HyperNetwork Model",
"insightface": "InsightFace Model",
"ipadapters": "ControlNet IP-Adapter Model",
"ipadapters/xlabs": "IP-Adapter from XLabs-AI",
"LLM": "LLM (aka Large-Language Model) is folder mapped (1 folder per model), append '/*' in the map",
"loras": "LoRA (aka Low-Ranking Adaption) Model",
"loras/xlabs": "LoRA Model from XLabs-AI",
"loras/flux": "LoRA Model trained on Flux.1 Dev or Flux.1 Schnell",
"reactor": "Reactor Model",
"reactor/faces": "Reactor Face Model",
"unet": "UNET Model Checkpoint (need separate CLIP and VAE Models)",
"upscale_models": "Upscaling Model (based on ESRGAN)",
"vae": "VAE En-/Decoder Model",
"vae-approx": "Approximate VAE Model"
}
# helper function called by "inline"-main() and ensure_shared_models_folders()
def init_shared_models_folders(send_SocketMessage:bool=True):
global SHARED_MODEL_FOLDERS
init_global_dict_from_file(SHARED_MODEL_FOLDERS, SHARED_MODEL_FOLDERS_FILE, "SHARED_MODEL_FOLDERS")
if os.path.exists(SHARED_MODEL_FOLDERS_FILE) and send_SocketMessage:
send_websocket_message('extend_ui_helper', {
'cmd': 'refreshModelTypes',
'message': 'New ModelTypes are available'
} )
return
### "inline"-main() ###
# init the SHARED_MODEL_FOLDERS
init_shared_models_folders(False) # dont send a WS-Message for "Model Downloader" at module init, to init/refresh its modelType list
# ----------
# helper function called from "app.py" via WebUI "Create Shared Folders" button on "Settings" tab
# ensures 'model_type' sub-folders for Model Mapping and the "Model Downloader" exists
# in the SHARED_MODELS_DIR (uses above initialized 'SHARED_MODEL_FOLDERS' dict)
def ensure_shared_models_folders():
try:
# init global module 'SHARED_MODEL_FOLDERS' dict: { 'model_type' (=subdir_names): 'app_model_dir'
# from app code or from external JSON 'SHARED_MODEL_FOLDERS_FILE' file
init_shared_models_folders(False) # (re-)read the SHARED_MODEL_FOLDERS_FILE again, if changed, but don't refresh modelTypes in "Model Downloader" yet
print(f"(re-)creating 'shared_models' model type sub-folders for Apps and the 'Model Downloader' in folder '{SHARED_MODELS_DIR}':")
# create the shared_models directory, if it doesn't exist yet
os.makedirs(f"{SHARED_MODELS_DIR}/", exist_ok=True) # append slash to make sure folder is created
# create a "__README.txt" file in the shared_models directory
readme_path = os.path.join(SHARED_MODELS_DIR, '__README.txt')
with open(readme_path, 'w') as readme_file:
readme_file.write("Upload your models to the appropriate folders:\n\n")
for model_type, model_type_description in SHARED_MODEL_FOLDERS.items():
shared_model_folderpath = os.path.join(SHARED_MODELS_DIR, model_type)
os.makedirs(os.path.dirname(f"{shared_model_folderpath}/"), exist_ok=True) # append trailing "/" to make sure the last sub-folder is created
print(f"'{model_type}' Folder created for '{model_type_description}'")
model_type_name = model_type
if model_type_name.endswith('s'): # model_type uses "plural" form with trailing "s"
model_type_name = model_type[:-1] # cut the last trailing 's'
readme_model_type_filename = os.path.join(shared_model_folderpath, f"{README_FILE_PREFIX}{model_type.replace('/', '-')}.txt") # translate "/" from grouping map rule into valid readme filename, e.g. "loras/flux" into "loras-flux"
readme_model_type_file = open(readme_model_type_filename, 'w')
readme_model_type_file.writelines(f"Put your '{model_type_name}' type models here, {model_type_description}")
readme_model_type_file.close()
readme_file.write(f"- {model_type}: for {model_type_name} models, {model_type_description}\n")
readme_file.write("\nThese models will be automatically linked to all supported apps.\n\n")
readme_file.write("Models directly downloaded into an app model folder will be\n")
readme_file.write("automatically pulled back into the corresponding shared folder and relinked back!\n")
# send a message for the "Model Downloader" to "refresh" its 'modelType' list
send_websocket_message('extend_ui_helper', {
'cmd': 'refreshModelTypes',
'message': 'New ModelTypes are available'
} )
return jsonify({'status': 'success', 'message': 'Shared model folders created successfully.'})
except Exception as e:
print(f"ERROR in shared_models:ensure_shared_model_folder() - Exception:\n{str(e)}")
return jsonify({'status': 'error', 'message': str(e)})
# ----------
# the below APP_INSTALL_DIRS_FILE will be read and used (if exists),
# otherwise this file will be generated with the content of the below default APP_INSTALL_DIRS dict
# this dict is very important, as it "defines" part of the symlink path,
# together with below defined SHARED_MODEL_APP_MAP (which uses relative path to the "app_install_dir" aka 'app_path')
APP_INSTALL_DIRS_FILE = f"{SHARED_MODELS_DIR}/{HIDDEN_FILE_PREFIX}_app_install_dirs.json"
APP_INSTALL_DIRS = {
# "app_name": "app_install_dir"
"A1111": "/workspace/stable-diffusion-webui",
"Forge": "/workspace/stable-diffusion-webui-forge",
"ComfyUI": "/workspace/ComfyUI",
"Kohya_ss": "/workspace/Kohya_ss",
"CUSTOM1": "/workspace/joy-caption-batch"
}
# the code from Madiator also defines a similar 'app_configs' dictionary
# but the idea here is also to allow "CUSTOM" Apps, installed by the user manually,
# to participate in "shared_models" model sharing
# app_configs = {
# 'bcomfy': {
# 'name': 'Better Comfy UI',
# 'command': 'cd /workspace/bcomfy && . ./bin/activate && cd /workspace/ComfyUI && python main.py --listen --port 3000 --enable-cors-header',
# 'venv_path': '/workspace/bcomfy',
# 'app_path': '/workspace/ComfyUI',
# 'port': 3000,
# },
# 'bforge': {
# 'name': 'Better Forge',
# 'command': 'cd /workspace/bforge && . ./bin/activate && cd /workspace/stable-diffusion-webui-forge && ./webui.sh -f --listen --enable-insecure-extension-access --api --port 7862',
# 'venv_path': '/workspace/bforge',
# 'app_path': '/workspace/stable-diffusion-webui-forge',
# 'port': 7862,
# },
# 'ba1111': {
# 'name': 'Better A1111',
# 'command': 'cd /workspace/ba1111 && . ./bin/activate && cd /workspace/stable-diffusion-webui && ./webui.sh -f --listen --enable-insecure-extension-access --api --port 7863',
# 'venv_path': '/workspace/ba1111',
# 'app_path': '/workspace/stable-diffusion-webui',
# 'port': 7863,
# }
# }
# MAP between Madiator's "app_configs" dict and the "APP_INSTALL_DIRS" dict used in this module
# TODO: this is temporary and should be merged/integrated better later
MAP_APPS = {
"bcomfy": "ComfyUI",
"bforge": "Forge",
"ba1111": "A1111"
}
# helper function called by main(), uses above "MAP_APPS" dict
# TODO: this is temporary and should be merged/integrated better later
def sync_with_app_configs_install_dirs():
print(f"Syncing 'app_configs' dict 'app_path' into the 'APP_INSTALL_DIRS' dict ...")
app_configs = get_app_configs()
for bapp_name, config in app_configs.items():
if bapp_name in MAP_APPS:
# get/sync the bapp_path from app_configs dict
bapp_path = app_configs[bapp_name]["app_path"]
print(f"\tSyncing 'app_path': '{bapp_path}' from app_configs for app 'name': '{bapp_name}'" )
APP_INSTALL_DIRS[MAP_APPS[bapp_name]] = bapp_path # update path in APP_INSTALL_DIRS
# show final synced APP_INSTALL_DIRS
print(f"\nUsing synched 'APP_INSTALL_DIRS':\n{PrettyDICT(APP_INSTALL_DIRS)}")
# init global module 'APP_INSTALL_DIRS' dict: { 'app_name': 'app_installdir' }
# default mapping from app code or (if exists) from external JSON 'APP_INSTALL_DIRS_FILE' file
# NOTE: this APP_INSTALL_DIRS_FILE is temporary synced with the app_configs dict
def init_app_install_dirs():
global APP_INSTALL_DIRS
init_global_dict_from_file(APP_INSTALL_DIRS, APP_INSTALL_DIRS_FILE, "APP_INSTALL_DIRS")
return
### "inline"-main() ###
# init the APP_INSTALL_DIRS and sync it from "app_configs" dict
init_app_install_dirs()
sync_with_app_configs_install_dirs() # TODO: this is temporary and should be merged/integrated better later
# ----------
SHARED_MODEL_APP_MAP_FILE = f"{SHARED_MODELS_DIR}/{HIDDEN_FILE_PREFIX}_shared_model_app_map.json"
# The dictionary 'model_type' "key" is relative to the SHARED_MODELS_DIR "/workspace/shared_models/" main folder.
# The sub dictionary 'app_model_folderpath' value is relative to the 'app_install_dir' value
# of the above APP_INSTALL_DIRS dictionary
# here is a list of all "known" model type dirs, and if they are used here (mapped) or
# if they are currently "unmapped":
#
# "Kohya_ss" (mapped): "/models"
# "ComfyUI" (mapped): "/models/checkpoints", "/models/clip", "/models/controlnet", "/models/embeddings", "/models/hypernetworks", "/models/ipadapter/"(???), "/models/loras", "/models/reactor"(???), "/models/unet", "/models/upscale_models", "/models/vae", "/models/vae_approx"
# "ComfyUI" (unmapped): "/models/clip_vision", "/models/diffusers", "/models/diffusion_models", "/models/gligen", "/models/photomaker", "/moedls/style_models",
# "A1111"/"Forge" (mapped): "/embeddings", "/models/ControlNet", "/models/ESRGAN", "/models/hypernetworks", "/models/insightface"(???), "/models/Lora", "/models/reactor", "/models/Stable-diffusion", "/models/text_encoder"(???), "/models/VAE", "/models/VAE-approx"
# "A1111"/"Forge" (unmapped): "/model/adetailer", "/models/BLIP", "/models/Codeformer", "models/deepbooru", "/model/Deforum", "/models/GFPGAN", "/models/karlo", "/models/Unet-onnx", "/models/Unet-trt"
SHARED_MODEL_APP_MAP = {
"ckpt": { # "model_type" (=subdir_name of SHARED_MODELS_DIR)
# "app_name": "app_model_folderpath" (for this "model_type", path is RELATIVE to "app_install_dir" of APP_INSTALL_DIRS map)
"ComfyUI": "/models/checkpoints",
"A1111": "/models/Stable-diffusion",
"Forge": "/models/Stable-diffusion",
"Kohya_ss": "/models" # flatten all "ckpt" / "unet" models here
},
"clip": {
"ComfyUI": "/models/clip",
"A1111": "/models/text_encoder",
"Forge": "/models/text_encoder"
},
"controlnet": {
"ComfyUI": "/models/controlnet",
"A1111": "/models/ControlNet",
"Forge": "/models/ControlNet"
#"A1111": "/extensions/sd-webui-controlnet/models", # SD1.5 ControlNets
#"Forge": "/extensions/sd-webui-controlnet/models" # SD1.5 ControlNets
},
# EMBEDDINGS map outside of models folder for FORGE / A1111
"embeddings": {
"ComfyUI": "/models/embeddings",
"A1111": "/embeddings",
"Forge": "/embeddings"
},
"hypernetworks": {
"ComfyUI": "/models/hypernetworks",
"A1111": "/models/hypernetworks",
"Forge": "/models/hypernetworks"
},
"insightface": {
"ComfyUI": "/models/insightface",
"A1111": "/models/insightface", # unverified location
"Forge": "/models/insightface" # unverified location
},
"ipadapters": {
"ComfyUI": "/models/ipadapter/",
"A1111": "/extensions/sd-webui-controlnet/models", # unverified location
"Forge": "/extensions/sd-webui-controlnet/models" # unverified location
},
"ipadapters/xlabs": { # sub-folders for XLabs-AI IP-Adapters
"ComfyUI": "/models/xlabs/ipadapters",
"A1111": "/extensions/sd-webui-controlnet/models", # flatten all "xlabs" ipadapters here
"Forge": "/extensions/sd-webui-controlnet/models" # flatten all "xlabs" ipadapters here
},
# some LoRAs get stored here in sub-folders, e.g. "/xlabs/*"
"loras": {
"ComfyUI": "/models/loras",
"A1111": "/models/Lora",
"Forge": "/models/Lora"
},
# Support "XLabs-AI" LoRA models
"loras/xlabs": { # special syntax for "grouping"
"ComfyUI": "/models/loras/xlabs",
"A1111": "/models/Lora", # flatten all "xlabs" LoRAs here
"Forge": "/models/Lora" # flatten all "xlabs" LoRAs here
},
# Support "Grouping" all FLUX LoRA models into a LoRA "flux" sub-folder for ComfyUI,
# which again need to be flattened for other apps
"loras/flux": {
"ComfyUI": "/models/loras/flux",
"A1111": "/models/Lora", # flatten all "flux" LoRAs here
"Forge": "/models/Lora" # flatten all "flux" LoRAs here
},
"reactor": {
"ComfyUI": "/models/reactor", # unverified location
"A1111": "/models/reactor",
"Forge": "/models/reactor",
},
"reactor/faces": {
"ComfyUI": "/models/reactor/faces", # unverified location
"A1111": "/models/reactor",
"Forge": "/models/reactor",
},
# UNET models map into the CKPT folders of all other apps, except for ComfyUI
"unet": {
"ComfyUI": "/models/unet",
"A1111": "/models/Stable-diffusion", # flatten all "ckpts" / "unet" models here
"Forge": "/models/Stable-diffusion", # flatten all "ckpts" / "unet" models here
"Kohya_ss": "/models" # flatten all "ckpt" / "unet" models here
},
"upscale_models": {
"ComfyUI": "/models/upscale_models",
"A1111": "/models/ESRGAN",
"Forge": "/models/ESRGAN"
},
"vae": {
"ComfyUI": "/models/vae",
"A1111": "/models/VAE",
"Forge": "/models/VAE"
},
"vae-approx": {
"ComfyUI": "/models/vae_approx",
"A1111": "/models/VAE-approx",
"Forge": "/models/VAE-approx"
},
# E.g. Custom Apps support for Joytag-Caption-Batch Tool (which uses the "Meta-Llama-3.1-8B" LLM)
# to share the model with e.g. ComfyUI. This LLM model come as full folders with more than one file!
# Pay attention to the special syntax for folder mappings (add a "/*" suffix to denote a folder mapping)
"LLM/Meta-Llama-3.1-8B/*": { # special syntax for "folder" symlink (the "/*" is mandatory)
"ComfyUI": "/models/LLM/Meta-Llama-3.1-8B/*", # special syntax for "folder" symlink, the "/*" is optional
"CUSTOM1": "/model/*" # special syntax for "folder" symlink, the "/*" is optional
}
}
# the "init_shared_model_app_map()" function initializes the
# global module 'SHARED_MODEL_APP_MAP' dict: 'model_type' -> 'app_name:app_model_dir' (relative path)
# which does a default mapping from app code or (if exists) from external JSON 'SHARED_MODEL_APP_MAP_FILE' file
def init_shared_model_app_map():
global SHARED_MODEL_APP_MAP
init_global_dict_from_file(SHARED_MODEL_APP_MAP, SHARED_MODEL_APP_MAP_FILE, "SHARED_MODEL_APP_MAP")
return
### "inline"-main() ###
# init the SHARED_MODEL_APP_MAP
init_shared_model_app_map()
# ----------
# helper function called by update_model_symlinks()
def remove_broken_model_symlinks(shared_model_folderpath:str, app_model_folderpath:str, model_type:str) -> int:
# process all files in app_model_folderpath
print(f"-> process broken '{model_type}' app_model file symlinks, which where removed from their corresponding shared_models sub-folder ...")
broken_modellinks_count = 0
broken_modellinks_info = ""
for app_model_filename in os.listdir(app_model_folderpath):
app_model_filepath = os.path.join(os.path.join(app_model_folderpath, app_model_filename))
# check for stale/broken model filelinks and folderlinks (LLMs)
if os.path.islink(app_model_filepath) and not os.path.exists(app_model_filepath):
# Remove existing stale/broken symlink
broken_modellinks_count = broken_modellinks_count + 1
dateInfo = "{:%B %d, %Y, %H:%M:%S GMT}".format(datetime.datetime.now())
broken_modellinks_info += f"\t{app_model_filename}\t[@ {dateInfo}]\n"
os.unlink(app_model_filepath) # try to unlink the file/folder symlink
# that normally is enough to remove the broken link (and the below code may never run)
# re-check if file/folder still exists
if os.path.exists(app_model_filepath): # if file/folder link still exists
if os.path.isdir(app_model_filepath): # broken folder link
shutil.rmtree(app_model_filepath) # remove the linked folder
else: # broken file link
os.remove(app_model_filepath) # remove the file link
print(f"\tremoved broken symlink for model '{app_model_filepath}'")
if broken_modellinks_count > 0:
# maintain (create/append to) a readme file for the app_model_folderpath target folder about the removed/deleted Model File Symlinks
readme_brokenlinks_models_filepath = os.path.join(app_model_folderpath, f"{README_FILE_PREFIX}brokenlinks-{model_type.replace('/', '-')}.txt") # translate "/" from grouping map rule into valid readme filename, e.g. "loras/flux" into "loras-flux"
if not os.path.exists(readme_brokenlinks_models_filepath): # no such readme file exists, so create it
fileHeader = f"Following broken model file links have been found and where deleted from this directory:\n\n"
file = open(readme_brokenlinks_models_filepath, 'w') # create the file
file.writelines(fileHeader) # and write the fileHeader once
file.writelines(broken_modellinks_info) # and add the broken Model File Links
file.close()
else: # readme file already existed from before
file = open(readme_brokenlinks_models_filepath, 'a') # append to file
file.writelines(broken_modellinks_info) # and add the broken Model File Links
file.close()
return broken_modellinks_count
# helper function called by update_model_symlinks()
def pull_unlinked_models_back_as_shared_models(shared_model_folderpath:str, app_model_folderpath:str, model_type:str) -> int:
# process all files in app_model_folderpath
print(f"-> process for possibly pulling-back '{model_type}' local app_model files into their corresponding shared_models sub-folder ...")
pulled_model_files_count = 0
pulled_model_files_info = ""
for app_model_filename in os.listdir(app_model_folderpath):
if app_model_filename.startswith(".") or app_model_filename.startswith(README_FILE_PREFIX):
continue # skip hidden filenames like ".DS_Store" (on macOS), ".keep" (on GitHub) and all "{README_FILE_PREFIX}*.txt" files
app_model_filepath = os.path.join(app_model_folderpath, app_model_filename)
if os.path.islink(app_model_filepath) or os.path.isdir(app_model_filepath):
continue # skip all already symlinked model files and sub-folders
# real file, potentially a model file which can be pulled back "home"
pulled_model_files_count = pulled_model_files_count + 1
print(f"processing app model '{app_model_filename}' ...")
shared_model_filepath = os.path.join(shared_model_folderpath, app_model_filename)
print(f"moving the file '{app_model_filename}' back to the '{model_type}' shared_models folder")
shutil.move(app_model_filepath, shared_model_filepath) # move it back to the shared_models model type folder
print(f"\tpulled-back local model '{app_model_filepath}'")
dateInfo = "{:%B %d, %Y, %H:%M:%S GMT}".format(datetime.datetime.now())
pulled_model_files_info += f"\t{app_model_filename}\t[@ {dateInfo}]\n"
### and re-link it back to this folder where it got just pulled back
# get the full path from shared model filename
src_filepath = os.path.join(shared_model_folderpath, app_model_filename)
dst_filepath = os.path.join(app_model_folderpath, app_model_filename)
if os.path.isfile(src_filepath) and not os.path.exists(dst_filepath):
os.symlink(src_filepath, dst_filepath)
print(f"\tre-created symlink {app_model_filename} -> {src_filepath} for pulled model")
if pulled_model_files_count > 0:
# maintain (create/append to) a readme file for the app_model_folderpath target folder about the pulled Model Files
readme_pulled_models_filepath = os.path.join(app_model_folderpath, f"{README_FILE_PREFIX}pulled-{model_type.replace('/', '-')}.txt") # translate "/" from grouping map rule into valid readme filename, e.g. "loras/flux" into "loras-flux"
if not os.path.exists(readme_pulled_models_filepath): # no such readme file exists, so create it
fileHeader = f"Following model files have been pulled from this directory into the shared_models directory '{shared_model_folderpath}' and re-linked here:\n\n"
file = open(readme_pulled_models_filepath, 'w') # create the file
file.writelines(fileHeader) # and write the fileHeader once
file.writelines(pulled_model_files_info) # and add the pulled Model Files
file.close()
else: # readme file already existed from before
file = open(readme_pulled_models_filepath, 'a') # append to file
file.writelines(pulled_model_files_info) # and add the pulled Model Files
file.close()
return pulled_model_files_count
# helper function called by update_model_symlinks()
def create_model_symlinks(shared_model_folderpath:str, app_model_folderpath:str, model_type:str) -> int:
# process all files in shared_model_folderpath
print(f"-> process for creating '{model_type}' app_model file symlinks from their corresponding shared_models sub-folder ...")
file_symlinks_created_count = 0
for shared_model_filename in os.listdir(shared_model_folderpath):
if shared_model_filename.startswith("."):
continue # skip hidden filenames like ".DS_Store" (on macOS), ".keep" (on GitHub)
# change the "readme-*.txt" files for the symlinked app folder models
if shared_model_filename.startswith(README_FILE_PREFIX):
# create a new readme file for the app_model_folderpath target folder
readme_synched_filename = os.path.join(app_model_folderpath, shared_model_filename.replace(README_FILE_PREFIX, f'{README_FILE_PREFIX}synced-'))
os.makedirs(os.path.dirname(readme_synched_filename), exist_ok=True) # ensure parent directory exists
file = open(readme_synched_filename, 'w')
file.writelines(f"This folder is synced from the shared_models '{model_type}' models type sub-folder at '{shared_model_folderpath}'.\n\nConsider to put such models there to share them across apps, instead of putting them here!")
file.close()
continue # skip the original "{README_FILE_PREFIX}*.txt" file
print(f"\tprocessing shared '{model_type}' model '{shared_model_filename}' ...")
# get the full path from shared model filename
src_filepath = os.path.join(shared_model_folderpath, shared_model_filename)
dst_filepath = os.path.join(app_model_folderpath, shared_model_filename) # the dst_filepath always has the SAME filename as the src_filepath
if not os.path.isfile(src_filepath): # srcFile is a sub-folder (e.g. "xlabs", or "flux")
# skip sub-folders, as these require a separate mapping rule to support "flattening" such models
# for apps which don't find their model_type models in sub-folders
# add a "model map" for "loras/flux" like the following:
# Support "Grouping" all FLUX LoRA models into a LoRA "flux" sub-folder for ComfyUI,
# which again need to be flattened for other apps
# "loras/flux": {
# "ComfyUI": "/models/loras/flux",
# "A1111": "/models/Lora", # flatten all "flux" LoRAs here
# "Forge": "/models/Lora" # flatten all "flux" LoRAs here
# }
print(f"\tthis is a sub-folder which should be mapped with a 'grouping' rule,\n\te.g. ""{model_type}/{shared_model_filename}: { ... }"" in '{SHARED_MODEL_APP_MAP_FILE}'")
continue
# create dst_filepath dirs for the parent folder, if needed
os.makedirs(os.path.dirname(dst_filepath), exist_ok=True)
if os.path.isfile(src_filepath) and not os.path.exists(dst_filepath):
os.symlink(src_filepath, dst_filepath)
print(f"\tcreated symlink {shared_model_filename} -> {src_filepath}")
file_symlinks_created_count = file_symlinks_created_count + 1
return file_symlinks_created_count
# helper function called from "app.py" via WebUI
#
# this is the main WORKER function running every 5 minutes
# or "on-demand" by the user from the WebUI via app.py:recreate_symlinks_route()->recreate_symlinks()
#
# this function uses following global module vars:
#
# README_FILE_PREFIX (str): "_readme-" <- README files are put in shared model type dirs and also in the app model type dirs,
# e.g. "_readme-*.txt", "_readme-synced-*.txt", "_readme-pulled-*.txt", "_readme-brokenlinks-*.txt",
# the "*" is filled with the dir-name of the corresponding model type subfolder of the SHARED_MODELS_DIR (str)
#
# DISABLE_PULLBACK_MODELS (bool) <- set via ENV (True, or False if not present [default])
# LOCAL_DEBUG (bool) <- set via ENV (True, or False if not present [default])
# MAKE_MAPPING_FILES_HIDDEN (bool): default=False (currently only controlled by app code)
#
# SHARED_MODELS_DIR (str): "/workspace/shared_models"
#
# SHARED_MODEL_FOLDERS_FILE (str): "_shared_models_folders.json" (based in SHARED_MODELS_DIR)
# SHARED_MODEL_FOLDERS (dict) <- init from code, then write/read from path SHARED_MODEL_FOLDERS_FILE
#
# APP_INSTALL_DIRS_FILE (str): "_app_install_dirs.json" (based in SHARED_MODELS_DIR)
# APP_INSTALL_DIRS (dict) <- init from code, then write/read from path SHARED_MODEL_FOLDERS_FILE,
# -> synced with global "app_configs" dict for 'app_path' with the use of MAP_APPS (dict)
# MAP_APPS (dict) <- used for mapping "app_configs" (dict) with APP_INSTALL_DIRS (dict)
#
# SHARED_MODEL_APP_MAP_FILE (str): "_shared_models_map.json" (based in SHARED_MODELS_DIR)
# SHARED_MODEL_APP_MAP (dict) <- init from code, then write/read from path SHARED_MODEL_FOLDERS_FILE
def update_model_symlinks() -> dict:
try:
print(f"Processing the master SHARED_MODELS_DIR: {SHARED_MODELS_DIR}")
if not os.path.exists(SHARED_MODELS_DIR):
message = f"Folder '{SHARED_MODELS_DIR}' does not exist, please create it first!"
return jsonify({'status': 'error', 'message': message})
file_model_symlinks_created_count = 0 # file model symlinks created
folder_model_symlinks_created_count = 0 # folder model symlinks created
broken_model_symlinks_count = 0 # broken symlinks to model files and folders
# "pull-back" model files can be disabled with ENV var "DISABLE_PULLBACK_MODELS=True"
pulled_model_files_count = 0 # pulled back model files (we not pull back folder models)
for model_type in SHARED_MODEL_APP_MAP:
print(f"\n### processing shared '{model_type}' model symlinks for all installed apps ...")
# check for special LLM folder symlink syntax
if not model_type.endswith("/*"): # normal file symlink in regular app_model_folderpath folder
create_folder_symlink = False
shared_model_folderpath = os.path.join(SHARED_MODELS_DIR, model_type)
else: # special case for folder symlink
create_folder_symlink = True
# strip the "/*" from the model_type (to deal with real folder names), before generating model_folderpaths
model_type_dirname = model_type.strip("/*")
shared_model_folderpath = os.path.join(SHARED_MODELS_DIR, model_type_dirname)
if not os.path.isdir(shared_model_folderpath):
print(f"shared_model_folderpath '{model_type}' does not exist, skipping")
continue # skipping non-existant shared_model_folderpath SRC folders
for app_name, app_install_dir in APP_INSTALL_DIRS.items():
if not os.path.exists(app_install_dir): # app is NOT installed
print(f"\n## app '{app_name}' is not installed, skipping")
continue # skipping non-installed app_install_dir for this model_type
print(f"\n## processing for app '{app_name}' ...")
if not (app_name in SHARED_MODEL_APP_MAP[model_type]):
print(f"-> there are no '{model_type}' symlink mappings defined for app '{app_name}', skipping")
continue # skipping non-existent app_name mapping for this model_type
app_model_folderpath = APP_INSTALL_DIRS[app_name] + SHARED_MODEL_APP_MAP[model_type][app_name]
print(f"# Processing the app's '{model_type}' folder '{app_model_folderpath}' ...")
if not create_folder_symlink: # normal file symlink in regular app_model_folderpath folder
# create the app model_type directory, if it doesn't exist
os.makedirs(f"{app_model_folderpath}/", exist_ok=True) # append slash to make sure folder is created
# first remove all broken/stale links
broken_model_type_symlinks_count = remove_broken_model_symlinks(shared_model_folderpath, app_model_folderpath, model_type)
if broken_model_type_symlinks_count > 0:
readme_filename = f"{README_FILE_PREFIX}brokenlinks-{model_type.replace('/', '-')}.txt" # translate "/" from grouping map rule into valid readme filename, e.g. "loras/flux" into "loras-flux"
print(f"-> found and removed #{broken_model_type_symlinks_count} broken link(s) for model type '{model_type}'\nfor more info about which model files symlinks were removed, look into the '{readme_filename}'")
# add them to its global counter
broken_model_symlinks_count = broken_model_symlinks_count + broken_model_type_symlinks_count
if not DISABLE_PULLBACK_MODELS:
# then try to pull back local, unlinked app models of this model type (they also get re-shared instantly back)
pulled_model_type_files_count = pull_unlinked_models_back_as_shared_models(shared_model_folderpath, app_model_folderpath, model_type)
if pulled_model_type_files_count > 0:
readme_filename = f"{README_FILE_PREFIX}pulled-{model_type.replace('/', '-')}.txt" # translate "/" from grouping map rule into valid readme filename, e.g. "loras/flux" into "loras-flux"
print(f"-> found and pulled back #{pulled_model_type_files_count} '{model_type}' model(s) into the corresponding shared_models sub-folder,\nfor more info about which model files were pulled, look into the '{readme_filename}'")
# add them to its global counter
pulled_model_files_count = pulled_model_files_count + pulled_model_type_files_count
# now share (symlink) all models of this type to app model path
file_model_type_symlinks_created_count = create_model_symlinks(shared_model_folderpath, app_model_folderpath, model_type)
if file_model_type_symlinks_created_count > 0:
# no readme details are tracked about created file symlinks
print(f"-> created #{file_model_type_symlinks_created_count} model file symlinks for '{model_type}' model(s)\nyou can see which models are now available in the app's '{app_model_folderpath}'")
# add them to its global counter
file_model_symlinks_created_count = file_model_symlinks_created_count + file_model_type_symlinks_created_count
else: # special case for folder symlink with LLM models, which install as folder
# e.g. app_model_folderpath = "/workspace/ComfyUI/models/LLM/Meta-Llama-3.1-8B/*"
# normally the target mapped/symlinked folder should also use a "/*" suffix,
# but that is not stricly required as we can handle that
# strip the "/*" from the model-map, to create the a "real" target folder per folder symlink
app_model_folderpath = app_model_folderpath.strip("/*")
# e.g. app_model_folderpath = "/workspace/ComfyUI/models/LLM/Meta-Llama-3.1-8B"
app_model_parent_dir, app_model_foldername = os.path.split(app_model_folderpath)
# e.g. app_model_parent_dir = "/workspace/ComfyUI/models/LLM"
os.makedirs(os.path.dirname(f"{app_model_parent_dir}/"), exist_ok=True) # append trailing "/" to make sure the last sub-folder is created
if os.path.exists(shared_model_folderpath) and not os.path.isfile(shared_model_folderpath) and not os.path.exists(app_model_folderpath):
# create a folder symlink
os.symlink(shared_model_folderpath, app_model_folderpath, target_is_directory=True)
# no readme details are tracked about created folder symlinks
print(f"\tcreated a folder symlink {app_model_foldername} -> {shared_model_folderpath}")
# the model_type counter for folder models is always 1, as each folder model is its own model_type
folder_symlinks_created_count = 1
# add this one folder symlink to its global (folder) counter (one-by-one)
folder_model_symlinks_created_count = folder_model_symlinks_created_count + folder_symlinks_created_count
pulled_models_info = "No Pull-Back"
if not DISABLE_PULLBACK_MODELS: # only show pulled models info, if "pull-back" is not disabled
pulled_models_info = f"Pulled({pulled_model_files_count})"
message = f"Links managed:\nFile({file_model_symlinks_created_count}), Folder({folder_model_symlinks_created_count}), Fixed({broken_model_symlinks_count}), {pulled_models_info}"
print(f"\n\nFinished updating all model type symlinks into their defined app model type directories.")
print(message)
print(f"\nFor further customizatons following files were now generated:\n")
print(f"- README.md: provided help about the other 3 JSON files:\n")
print(f"- SHARED_MODEL_FOLDERS_FILE '{SHARED_MODEL_FOLDERS_FILE}':\nprovides examples of used 'model_type' directory names for different models.\n")
print(f"- APP_INSTALL_DIRS_FILE '{APP_INSTALL_DIRS_FILE}':\nprovides examples of used 'app_name': 'app_install_dir' mappings and is used together with the SHARED_MODEL_APP_MAP_FILE.\n")
print(f"- SHARED_MODEL_APP_MAP_FILE '{SHARED_MODEL_APP_MAP_FILE}':\nprovides examples of used 'model_type' -> 'app_model_dir' mappings.\n")
print("Model symlinks updated.")
return jsonify({'status': 'success', 'message': message}) # 'Symlinks (re-)created successfully.'
except Exception as e:
print(f"ERROR in shared_models:update_model_symlinks() - Exception:\n{str(e)}")
return jsonify({'status': 'error', 'message': str(e)})
# promote the README
print("To get started with all features of 'shared_models', consult the comprehensive README file")
print('\t"/app/tests/README-SHARED_MODELS.txt"\nIt comes with a Test script and Test data.\n')
print("TESTDATA AND EXPLANATION OF MAPPING EVERYTHING YOU WANT\n")
print('In the folder "/app/tests" you find the following files:')
print('\t- "README-SHARED_MODELS.txt" (this file)')
print('\t- "populate_testdata.sh" (bash script to un-tar and expand all testdata into the "/workspace" folder)')
print('\t- "testdata_shared_models_link.tar.gz" (Testcase #1, read below)')
print('\t- "testdata_stable-diffusion-webui_pull.tar.gz" (Testcase #2, read below)')
print('\t- "testdata_installed_apps_pull.tar.gz" (Testcase #3, read below)\n')

0
official-templates/better-ai-launcher/utils/ssh_utils.py → official-templates/better-ai-launcher/app/utils/ssh_utils.py
View file

0
official-templates/better-ai-launcher/utils/websocket_utils.py → official-templates/better-ai-launcher/app/utils/websocket_utils.py
View file

23
official-templates/better-ai-launcher/docker-compose.debug.yml Normal file
View file

@ -0,0 +1,23 @@
services:
mydev:
image: madiator2011/better-launcher:dev
build:
context: .
dockerfile: ./Dockerfile
command: ["sh", "-c", "pip install debugpy -t /tmp && python /tmp/debugpy --listen 0.0.0.0:5678 -m flask run --host 0.0.0.0 --port 7222"]
ports:
- 5678:5678 # (random) port for debuggy (adjust together in above "command")
- 7222:7222 # main Flask app port better-launcher "App-Manager"
- 8181:8181 # File-Browser
- 7777:7777 # VSCode-Server
env_file:
- .env # pass additional env-vars (hf_token, civitai token, ssh public-key) from ".env" file to container
environment:
- LOCAL_DEBUG=True # change app to localhost Urls and local Websockets (unsecured)
- FLASK_APP=app/app.py
- FLASK_ENV=development # changed from "production"
- GEVENT_SUPPORT=True # gevent monkey-patching is being used, enable gevent support in the debugger
#- "FLASK_DEBUG": "0" # "1" allows debugging in Chrome, but then VSCode debugger not works
volumes:
- ./app:/app:rw
- ${HOME}/Projects/Docker/madiator:/workspace:rw # TODO: create the below folder before you run!

16
official-templates/better-ai-launcher/docker-compose.yml Normal file
View file

@ -0,0 +1,16 @@
services:
mydev:
image: madiator2011/better-launcher:dev
build:
context: .
dockerfile: ./Dockerfile
env_file:
- .env # pass additional env-vars (hf_token, civitai token, ssh public-key) from ".env" file to container
environment:
- LOCAL_DEBUG=True # change app to localhost Urls and local Websockets (unsecured)
volumes:
- ${HOME}/Projects/Docker/madiator:/workspace:rw # # TODO: create the below folder before you run!
ports:
- 7222:7222 # main Flask app port better-launcher "App-Manager"
- 8181:8181 # File-Browser
- 7777:7777 # VSCode-Server

34
official-templates/better-ai-launcher/env.txt Normal file
View file

@ -0,0 +1,34 @@
### for local debugging fill-in/update YOUR env values and secrets
### and rename this file as ".env" (hidden file with NO extension)
### the ".env" file is also added to the ".dockerignore" and ".gitignore" files,
### to NOT pass them to the container (".dockerignore") or to GitHub (".gitignore")
###
### this ".env" file is then "passed" into the debugging container via ".vscode/tasks.json"
### in the "dockerRun" section:
#
# "dockerRun": {
# "envFiles": ["${workspaceFolder}/.env"], // pass additional env-vars from ".env" file to container
### Build Vars ###
IMAGE_BASE=madiator2011/better-launcher
IMAGE_TAG=dev
### APP specific Vars ###
DISABLE_PULLBACK_MODELS=False
# the default is, that app model files, which are found locally (in only one app),
# get automatically "pulled-back" into the '/workspace/shared_models' folder.
# From there they will be re-linked back not only to their own "pulled-back" model-type folder,
# but also will be linked back into all other corresponding app model-type folders.
# So the "pulled-back" model is automatically shared to all installed apps.
#
# if you NOT want this behaviour, then set DISABLE_PULLBACK_MODELS=True
### USER specific Vars and Secrets (Tokens) - TODO: adjust this for your personal settings ###
PUBLIC_KEY=ssh-ed25519 XXX...XXX usermail@domain.com
HF_TOKEN=hf_XXX...XXX
CIVITAI_API_TOKEN=XXX.XXX
### RUNPOD specific Vars ###
RUNPOD_PUBLIC_IP=127.0.0.1
RUNPOD_TCP_PORT_22=22

7
official-templates/better-ai-launcher/nginx/README.md Normal file
View file

@ -0,0 +1,7 @@
## Build Options
To build with default options, run `docker buildx bake`, to build a specific target, run `docker buildx bake <target>`.
## Ports
- 22/tcp (SSH)

0
official-templates/better-ai-launcher/nginx.conf → official-templates/better-ai-launcher/nginx/nginx.conf
View file

0
official-templates/better-ai-launcher/readme.html → official-templates/better-ai-launcher/nginx/readme.html
View file

255
official-templates/better-ai-launcher/utils/model_utils.py
View file

@ -1,255 +0,0 @@
import os
import requests
from urllib.parse import unquote, urlparse
from tqdm import tqdm
import json
import re
import time
import math
SHARED_MODELS_DIR = '/workspace/shared_models'
MODEL_TYPE_MAPPING = {
'Checkpoint': 'Stable-diffusion',
'LORA': 'Lora',
'LoCon': 'Lora',
'TextualInversion': 'embeddings',
'VAE': 'VAE',
'Hypernetwork': 'hypernetworks',
'AestheticGradient': 'aesthetic_embeddings',
'ControlNet': 'controlnet',
'Upscaler': 'ESRGAN'
}
def ensure_shared_folder_exists():
for folder in ['Stable-diffusion', 'Lora', 'embeddings', 'VAE', 'hypernetworks', 'aesthetic_embeddings', 'controlnet', 'ESRGAN']:
os.makedirs(os.path.join(SHARED_MODELS_DIR, folder), exist_ok=True)
def check_civitai_url(url):
prefix = "civitai.com"
try:
if prefix in url:
if "civitai.com/api/download" in url:
version_id = url.strip("/").split("/")[-1]
return False, True, None, int(version_id)
subpath = url[url.find(prefix) + len(prefix):].strip("/")
url_parts = subpath.split("?")
if len(url_parts) > 1:
model_id = url_parts[0].split("/")[1]
version_id = url_parts[1].split("=")[1]
return True, False, int(model_id), int(version_id)
else:
model_id = subpath.split("/")[1]
return True, False, int(model_id), None
except (ValueError, IndexError):
print("Error parsing Civitai model URL")
return False, False, None, None
def check_huggingface_url(url):
parsed_url = urlparse(url)
if parsed_url.netloc not in ["huggingface.co", "huggingface.com"]:
return False, None, None, None, None
path_parts = [p for p in parsed_url.path.split("/") if p]
if len(path_parts) < 5 or (path_parts[2] != "resolve" and path_parts[2] != "blob"):
return False, None, None, None, None
repo_id = f"{path_parts[0]}/{path_parts[1]}"
branch_name = path_parts[3]
remaining_path = "/".join(path_parts[4:])
folder_name = os.path.dirname(remaining_path) if "/" in remaining_path else None
filename = unquote(os.path.basename(remaining_path))
return True, repo_id, filename, folder_name, branch_name
def download_model(url, model_name, model_type, send_websocket_message, civitai_token=None, hf_token=None, version_id=None, file_index=None):
ensure_shared_folder_exists()
is_civitai, is_civitai_api, model_id, _ = check_civitai_url(url)
is_huggingface, repo_id, hf_filename, hf_folder_name, hf_branch_name = check_huggingface_url(url)
if is_civitai or is_civitai_api:
if not civitai_token:
return False, "Civitai token is required for downloading from Civitai"
success, message = download_civitai_model(url, model_name, model_type, send_websocket_message, civitai_token, version_id, file_index)
elif is_huggingface:
success, message = download_huggingface_model(url, model_name, model_type, send_websocket_message, repo_id, hf_filename, hf_folder_name, hf_branch_name, hf_token)
else:
return False, "Unsupported URL"
if success:
send_websocket_message('model_download_progress', {
'percentage': 100,
'stage': 'Complete',
'message': 'Download complete and symlinks updated'
})
return success, message
def download_civitai_model(url, model_name, model_type, send_websocket_message, civitai_token, version_id=None, file_index=None):
try:
is_civitai, is_civitai_api, model_id, url_version_id = check_civitai_url(url)
headers = {'Authorization': f'Bearer {civitai_token}'}
if is_civitai_api:
api_url = f"https://civitai.com/api/v1/model-versions/{url_version_id}"
else:
api_url = f"https://civitai.com/api/v1/models/{model_id}"
response = requests.get(api_url, headers=headers)
response.raise_for_status()
model_data = response.json()
if is_civitai_api:
version_data = model_data
model_data = version_data['model']
else:
if version_id:
version_data = next((v for v in model_data['modelVersions'] if v['id'] == version_id), None)
elif url_version_id:
version_data = next((v for v in model_data['modelVersions'] if v['id'] == url_version_id), None)
else:
version_data = model_data['modelVersions'][0]
if not version_data:
return False, f"Version ID {version_id or url_version_id} not found for this model."
civitai_model_type = model_data['type']
model_type = MODEL_TYPE_MAPPING.get(civitai_model_type, 'Stable-diffusion')
files = version_data['files']
if file_index is not None and 0 <= file_index < len(files):
file_to_download = files[file_index]
elif len(files) > 1:
# If there are multiple files and no specific file was chosen, ask the user to choose
file_options = [{'name': f['name'], 'size': f['sizeKB'], 'type': f['type']} for f in files]
return True, {
'choice_required': {
'type': 'file',
'model_id': model_id,
'version_id': version_data['id'],
'files': file_options
}
}
else:
file_to_download = files[0]
download_url = file_to_download['downloadUrl']
if not model_name:
model_name = file_to_download['name']
model_path = os.path.join(SHARED_MODELS_DIR, model_type, model_name)
os.makedirs(os.path.dirname(model_path), exist_ok=True)
return download_file(download_url, model_path, send_websocket_message, headers)
except requests.RequestException as e:
return False, f"Error downloading from Civitai: {str(e)}"
def download_huggingface_model(url, model_name, model_type, send_websocket_message, repo_id, hf_filename, hf_folder_name, hf_branch_name, hf_token=None):
try:
from huggingface_hub import hf_hub_download
if not model_name:
model_name = hf_filename
model_path = os.path.join(SHARED_MODELS_DIR, model_type, model_name)
os.makedirs(os.path.dirname(model_path), exist_ok=True)
send_websocket_message('model_download_progress', {
'percentage': 0,
'stage': 'Downloading',
'message': f'Starting download from Hugging Face: {repo_id}'
})
kwargs = {
'repo_id': repo_id,
'filename': hf_filename,
'subfolder': hf_folder_name,
'revision': hf_branch_name,
'local_dir': os.path.dirname(model_path),
'local_dir_use_symlinks': False
}
if hf_token:
kwargs['token'] = hf_token
local_file = hf_hub_download(**kwargs)
send_websocket_message('model_download_progress', {
'percentage': 100,
'stage': 'Complete',
'message': f'Download complete: {model_name}'
})
return True, f"Successfully downloaded {model_name} from Hugging Face"
except Exception as e:
return False, f"Error downloading from Hugging Face: {str(e)}"
def download_file(url, filepath, send_websocket_message, headers=None):
try:
response = requests.get(url, stream=True, headers=headers)
response.raise_for_status()
total_size = int(response.headers.get('content-length', 0))
block_size = 8192
downloaded_size = 0
start_time = time.time()
with open(filepath, 'wb') as file:
for data in response.iter_content(block_size):
size = file.write(data)
downloaded_size += size
current_time = time.time()
elapsed_time = current_time - start_time
if elapsed_time > 0:
speed = downloaded_size / elapsed_time
percentage = (downloaded_size / total_size) * 100 if total_size > 0 else 0
eta = (total_size - downloaded_size) / speed if speed > 0 else 0
send_websocket_message('model_download_progress', {
'percentage': round(percentage, 2),
'speed': f"{speed / (1024 * 1024):.2f} MB/s",
'eta': int(eta),
'stage': 'Downloading',
'message': f'Downloaded {format_size(downloaded_size)} / {format_size(total_size)}'
})
send_websocket_message('model_download_progress', {
'percentage': 100,
'stage': 'Complete',
'message': f'Download complete: {os.path.basename(filepath)}'
})
return True, f"Successfully downloaded {os.path.basename(filepath)}"
except requests.RequestException as e:
return False, f"Error downloading file: {str(e)}"
def get_civitai_file_size(url, token):
headers = {'Authorization': f'Bearer {token}'}
try:
response = requests.head(url, headers=headers, allow_redirects=True)
return int(response.headers.get('content-length', 0))
except:
return None
def get_huggingface_file_size(repo_id, filename, folder_name, branch_name, token):
from huggingface_hub import hf_hub_url, HfApi
try:
api = HfApi()
file_info = api.hf_hub_url(repo_id, filename, subfolder=folder_name, revision=branch_name)
response = requests.head(file_info, headers={'Authorization': f'Bearer {token}'} if token else None)
return int(response.headers.get('content-length', 0))
except:
return None
def format_size(size_in_bytes):
if size_in_bytes == 0:
return "0 B"
size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
i = int(math.floor(math.log(size_in_bytes, 1024)))
p = math.pow(1024, i)
s = round(size_in_bytes / p, 2)
return f"{s} {size_name[i]}"