Last large refactoring movement, now it's polish and getting things working...

Author: oxysoft

core/cmdargs.py

@@ -4,53 +4,57 @@
 from core.paths import *
  # TODO do not use cargs
 
+# The parsed args (global)
+cargs = None
+
 parser = argparse.ArgumentParser()
-parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))
-parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)
-parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")
-parser.add_argument("--no-half-vae", action='store_true', help="do not switch the VAE model to 16-bit floats")
-parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")
-parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")
-parser.add_argument("--embeddings-dir", type=str, default=rootdir / 'embeddings', help="embeddings directory for textual inversion (default: embeddings)")
-parser.add_argument("--hypernetwork-dir", type=str, default=modeldir / 'modules/hypernetworks', help="hypernetwork directory")
-parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")
-parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a little speed for low VRM usage")
-parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a lot of speed for very low VRM usage")
-parser.add_argument("--always-batch-cond-uncond", action='store_true', help="disables cond/uncond batching that is enabled to save memory with --medvram or --lowvram")
-parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.")
-parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")
-parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site (doesn't work for me but you might have better luck)")
-parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to gradio --share", default=None)
-parser.add_argument("--codeformer-models-path", type=str, help="Path to directory with codeformer model file(s).", default=modeldir / 'Codeformer')
-parser.add_argument("--gfpgan-models-path", type=str, help="Path to directory with GFPGAN model file(s).", default=(modeldir / 'GFPGAN').as_posix())
-parser.add_argument("--esrgan-models-path", type=str, help="Path to directory with ESRGAN model file(s).", default=(modeldir / 'ESRGAN').as_posix())
-parser.add_argument("--bsrgan-models-path", type=str, help="Path to directory with BSRGAN model file(s).", default=(modeldir / 'BSRGAN').as_posix())
-parser.add_argument("--realesrgan-models-path", type=str, help="Path to directory with RealESRGAN model file(s).", default=(modeldir / 'RealESRGAN').as_posix())
-parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=(modeldir / 'ScuNET').as_posix())
-parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=modeldir / 'SwinIR')
-parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=modeldir / 'LDSR')
-parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")
-parser.add_argument("--force-enable-xformers", action='store_true', help="enable xformers for cross attention layers regardless of whether the checking code thinks you can run it; do not make bug reports if this fails to work")
-parser.add_argument("--opt-split-attention", action='store_true', help="force-enables Doggettx's cross-attention layer optimization. By default, it's on for torch cuda.")
-parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="force-enables InvokeAI's cross-attention layer optimization. By default, it's on when cuda is unavailable.")
-parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")
-parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")
-parser.add_argument("--use-cpu", nargs='+', choices=['StableDiffusion', 'GFPGAN', 'BSRGAN', 'ESRGAN', 'SCUNet', 'CodeFormer'], help="use CPU as torch device for specified modules", default=[])
-parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
-parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)
-parser.add_argument("--show-negative-prompt", action='store_true', help="does not do anything", default=False)
-parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=rootdir / 'ui-config.json')
-parser.add_argument("--hide-ui-dir-config", action='store_true', help="hide directory configuration from webui", default=False)
-parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=rootdir / 'config.json')
-parser.add_argument("--gradio-debug", action='store_true', help="launch gradio with --debug option")
-parser.add_argument("--gradio-auth", type=str, help='set gradio authentication like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)
-parser.add_argument("--gradio-img2img-tool", type=str, help='gradio image uploader tool: can be either editor for ctopping, or color-sketch for drawing', choices=["color-sketch", "editor"], default="editor")
-parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last")
-parser.add_argument("--styles-file", type=str, help="filename to use for styles", default=rootdir / 'styles.csv')
-parser.add_argument("--autolaunch", action='store_true', help="open the webui URL in the system's default browser upon launch", default=False)
-parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)
-parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=False)
-parser.add_argument('--vae-path', type=str, help='Path to Variational Autoencoders model', default=None)
-parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)
+parser.add_argument("--dry", type=str, help="Only install and test the core, do not launch server.")
+
+# parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))
+# parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)
+# parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")
+# parser.add_argument("--no-half-vae", action='store_true', help="do not switch the VAE model to 16-bit floats")
+# parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")
+# parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")
+# parser.add_argument("--embeddings-dir", type=str, default=rootdir / 'embeddings', help="embeddings directory for textual inversion (default: embeddings)")
+# parser.add_argument("--hypernetwork-dir", type=str, default=modeldir / 'modules/hypernetworks', help="hypernetwork directory")
+# parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")
+# parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a little speed for low VRM usage")
+# parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a lot of speed for very low VRM usage")
+# parser.add_argument("--always-batch-cond-uncond", action='store_true', help="disables cond/uncond batching that is enabled to save memory with --medvram or --lowvram")
+# parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.")
+# parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")
+# parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site (doesn't work for me but you might have better luck)")
+# parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to gradio --share", default=None)
+# parser.add_argument("--codeformer-models-path", type=str, help="Path to directory with codeformer model file(s).", default=modeldir / 'Codeformer')
+# parser.add_argument("--gfpgan-models-path", type=str, help="Path to directory with GFPGAN model file(s).", default=(modeldir / 'GFPGAN').as_posix())
+# parser.add_argument("--esrgan-models-path", type=str, help="Path to directory with ESRGAN model file(s).", default=(modeldir / 'ESRGAN').as_posix())
+# parser.add_argument("--bsrgan-models-path", type=str, help="Path to directory with BSRGAN model file(s).", default=(modeldir / 'BSRGAN').as_posix())
+# parser.add_argument("--realesrgan-models-path", type=str, help="Path to directory with RealESRGAN model file(s).", default=(modeldir / 'RealESRGAN').as_posix())
+# parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=(modeldir / 'ScuNET').as_posix())
+# parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=modeldir / 'SwinIR')
+# parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=modeldir / 'LDSR')
+# parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")
+# parser.add_argument("--force-enable-xformers", action='store_true', help="enable xformers for cross attention layers regardless of whether the checking code thinks you can run it; do not make bug reports if this fails to work")
+# parser.add_argument("--opt-split-attention", action='store_true', help="force-enables Doggettx's cross-attention layer optimization. By default, it's on for torch cuda.")
+# parser.add_argument("--opt-split-attention-invokeai", action='store_true', help="force-enables InvokeAI's cross-attention layer optimization. By default, it's on when cuda is unavailable.")
+# parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")
+# parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")
+# parser.add_argument("--use-cpu", nargs='+', choices=['stable_diffusion', 'GFPGAN', 'BSRGAN', 'ESRGAN', 'SCUNet', 'CodeFormer'], help="use CPU as torch device for specified modules", default=[])
+# parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
+# parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)
+# parser.add_argument("--show-negative-prompt", action='store_true', help="does not do anything", default=False)
+# parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=rootdir / 'ui-config.json')
+# parser.add_argument("--hide-ui-dir-config", action='store_true', help="hide directory configuration from webui", default=False)
+# parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=rootdir / 'config.json')
+# parser.add_argument("--gradio-debug", action='store_true', help="launch gradio with --debug option")
+# parser.add_argument("--gradio-auth", type=str, help='set gradio authentication like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)
+# parser.add_argument("--gradio-img2img-tool", type=str, help='gradio image uploader tool: can be either editor for ctopping, or color-sketch for drawing', choices=["color-sketch", "editor"], default="editor")
+# parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last")
+# parser.add_argument("--styles-file", type=str, help="filename to use for styles", default=rootdir / 'styles.csv')
+# parser.add_argument("--autolaunch", action='store_true', help="open the webui URL in the system's default browser upon launch", default=False)
+# parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)
+# parser.add_argument("--disable-console-progressbars", action='store_true', help="do not output progressbars to console", default=False)
+# parser.add_argument('--vae-path', type=str, help='Path to Variational Autoencoders model', default=None)
+# parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)
 
-cargs = None

core/devicelib.py

@@ -3,7 +3,6 @@
 import torch
 
 from core import printing
-from core.cmdargs import cargs
 
 
 def get_optimal_device():
@@ -50,12 +49,9 @@ def randn_without_seed(shape):
     return torch.randn(shape, device=device)
 
 
-def autocast(disable=False):
-    if disable:
-        return contextlib.nullcontext()
-
-    if dtype == torch.float32 or cargs.precision == "full":
-        return contextlib.nullcontext()
+def autocast(enable=True):
+    if not enable: return contextlib.nullcontext()
+    if dtype == torch.float32: return contextlib.nullcontext() # or cargs.precision == "full"
 
     return torch.autocast("cuda")
 
@@ -79,7 +75,7 @@ def autocast(disable=False):
 # device_codeformer = \
 #     (cpu if x in cargs.use_cpu else get_optimal_device()
 #      for x
-#      in ['StableDiffusion', 'GFPGAN', 'BSRGAN', 'ESRGAN', 'SCUNet', 'CodeFormer'])
+#      in ['stable_diffusion', 'GFPGAN', 'BSRGAN', 'ESRGAN', 'SCUNet', 'CodeFormer'])
 
 # device = device
 

core/installing.py

@@ -82,6 +82,8 @@ def git_clone(url, dir, name, commithash=None):
     if commithash is not None:
         run(f'"{git}" -C {dir} checkout {commithash}', None, "Couldn't checkout {name}'s hash: {commithash}")
 
+    sys.path.append(dir)
+
 
 def move_files(src_path: str, dest_path: str, ext_filter: str = None):
     try:

core/jobs.py

@@ -1,4 +1,5 @@
 import uuid
+from abc import abstractmethod
 from datetime import datetime
 
 from tqdm import tqdm
@@ -14,9 +15,16 @@ class JobParams:
     def __init__(self, **kwargs):
         self.__dict__.update(kwargs)
 
+    @abstractmethod
+    def get_plugin_impl(self):
+        pass
+
     def plugin(self):
         return self.job.plugin()
 
+    def on_start(self, job):
+        pass
+
 
 class Job:
     def __init__(self, plugin_id: str, plugin_func: str, parameters: JobParams):
@@ -88,6 +96,8 @@ def process(self, job):
         self.queued.remove(job)
         self.processing.append(job)
 
+        job.p.on_start(job)
+
         webui.emit('started_job', job.job_id)
 
     def cancel(self, job):
@@ -168,15 +178,22 @@ def clear(self):
             self.tqdm = None
 
 
-def get_job(job_id):
+def get(job_id):
+    """
+    Get a job by id.
+    The job must be queued or processing.
+    """
+    if isinstance(job_id, Job):
+        return job_id
+
     for job in queue.all:
         if job.job_id == job_id:
             return job
     raise Exception("Job not found")
 
 
 def finish(job_id):
-    job = get_job(job_id)
+    job = get(job_id)
 
     job.update(1)
     queue.finish(job)
@@ -201,10 +218,10 @@ def enqueue(job):
     queue.enqueue(job)
 
 
-def new_job(plugid, jobname, parameters):
+def new_job(plugid, name, jobparams):
     global total_created
     total_created += 1
-    j = Job(plugid, jobname, parameters)
+    j = Job(plugid, name, jobparams)
     enqueue(j)
 
     return j
@@ -214,3 +231,16 @@ def new_job(plugid, jobname, parameters):
 queue = JobQueue()
 total_created = 0
 tqdm_total = JobTQDM()
+
+
+def is_processing(id):
+    """
+    Check if a job is processing (by ID)
+    """
+    j = get(id)
+    if j is not None:
+        return j in queue.processing
+
+
+def is_any_processing():
+    return len(queue.processing) > 0

core/modellib.py

@@ -3,9 +3,9 @@
 from pathlib import Path
 from urllib.parse import urlparse
 
-from basicsr.utils.download_util import load_file_from_url
-
 import torch
+from torch.hub import download_url_to_file, get_dir
+
 from core.devicelib import get_optimal_device
 
 import pickle
@@ -114,6 +114,36 @@ def load_models(model_path: Path, model_url: Path = None, command_path: Path = N
 
     return output
 
+def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+    """Load file form http url, will download models if necessary.
+
+    Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
+
+    Args:
+        url (str): URL to be downloaded.
+        model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
+            Default: None.
+        progress (bool): Whether to show the download progress. Default: True.
+        file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
+
+    Returns:
+        str: The path to the downloaded file.
+    """
+    if model_dir is None:  # use the pytorch hub_dir
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, 'checkpoints')
+
+    os.makedirs(model_dir, exist_ok=True)
+
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if file_name is not None:
+        filename = file_name
+    cached_file = os.path.abspath(os.path.join(model_dir, filename))
+    if not os.path.exists(cached_file):
+        print(f'Downloading: "{url}" to {cached_file}\n')
+        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
+    return cached_file
 
 # def load_upscalers():
 #     sd = shared.script_path
@@ -217,11 +247,9 @@ def check_pt(filename):
 
 
 def load(filename, *args, **kwargs):
-    import shared
-
     try:
-        if not shared.cmd_opts.disable_safe_unpickle:
-            check_pt(filename)
+        # if not shared.cmd_opts.disable_safe_unpickle:
+        check_pt(filename)
 
     except Exception:
         print(f"Error verifying pickled file from {filename}:", file=sys.stderr)