Skip to content

Commit

Permalink
adds automatic downloads of AI models
Browse files Browse the repository at this point in the history
  • Loading branch information
zackees committed Dec 16, 2023
1 parent 7f171b1 commit 58f9416
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 1 deletion.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -127,3 +127,4 @@ dmypy.json

# Pyre type checker
.pyre/
AI/*
15 changes: 14 additions & 1 deletion RealScaler.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@

# Standard library imports
import os
import sys
from shutil import rmtree as remove_directory
from timeit import default_timer as timer
Expand Down Expand Up @@ -134,6 +135,8 @@
githubme = "https://github.com/Djdefrag/ReSRScaler"
telegramme = "https://linktr.ee/j3ngystudio"

# Remote URL where the AI *.pth assets can be downloaded from.
REMOTE_AI_URL = "https://zackees.github.io/ai-image-upscaler-pth"

# SRVGGNetCompact
SRVGGNetCompact_vram_multiplier = 1.7
Expand Down Expand Up @@ -190,6 +193,14 @@
]


def ensure_downloaded(modelfile: str, outpath: str) -> None:
"""Download a model file if it doesn't exist."""
if os.path.exists(outpath):
return # Already downloaded
asset_url = f"{REMOTE_AI_URL}/{modelfile}"
from download import download # Lazy import
download(asset_url, outpath)


# AI models -------------------

Expand All @@ -202,7 +213,9 @@ def load_AI_model(
) -> any:

write_process_status(processing_queue, f"Loading AI model")
model_path = find_by_relative_path(f"AI{os_separator}{selected_AI_model}.pth")
model_file = f"{selected_AI_model}.pth"
model_path = find_by_relative_path(f"AI{os_separator}{model_file}")
ensure_downloaded(model_file, model_path)

if selected_AI_model == 'RealESRGANx4':
model = RealESRGAN_Net()
Expand Down
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,4 @@ packaging
moviepy
opencv-python-headless
nuitka
download

0 comments on commit 58f9416

Please sign in to comment.