Merge branch 'master' of https://github.com/cluder/stable-diffusion-webui into 4448_fix_ckpt_cache

CODEOWNERS

 *       @AUTOMATIC1111
-/localizations/ar_AR.json   @xmodar @blackneoo
-/localizations/de_DE.json   @LunixWasTaken
-/localizations/es_ES.json   @innovaciones
-/localizations/fr_FR.json   @tumbly
-/localizations/it_IT.json   @EugenioBuffo
-/localizations/ja_JP.json   @yuuki76
-/localizations/ko_KR.json   @36DB
-/localizations/pt_BR.json   @M-art-ucci
-/localizations/ru_RU.json   @kabachuha
-/localizations/tr_TR.json   @camenduru
-/localizations/zh_CN.json   @dtlnor @bgluminous
-/localizations/zh_TW.json   @benlisquare
+
+# if you were managing a localization and were removed from this file, this is because
+# the intended way to do localizations now is via extensions. See:
+# https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Developing-extensions
+# Make a repo with your localization and since you are still listed as a collaborator
+# you can add it to the wiki page yourself. This change is because some people complained
+# the git commit log is cluttered with things unrelated to almost everyone and
+# because I believe this is the best overall for the project to handle localizations almost
+# entirely without my oversight.
+
+

javascript/ui.js

@@ -208,4 +208,6 @@ function update_token_counter(button_id) {
 function restart_reload(){
     document.body.innerHTML='<h1 style="font-family:monospace;margin-top:20%;color:lightgray;text-align:center;">Reloading...</h1>';
     setTimeout(function(){location.reload()},2000)
+
+    return []
 }

modules/api/api.py

@@ -63,6 +63,7 @@ class Api:
         self.app.add_api_route("/sdapi/v1/extra-batch-images", self.extras_batch_images_api, methods=["POST"], response_model=ExtrasBatchImagesResponse)
         self.app.add_api_route("/sdapi/v1/png-info", self.pnginfoapi, methods=["POST"], response_model=PNGInfoResponse)
         self.app.add_api_route("/sdapi/v1/progress", self.progressapi, methods=["GET"], response_model=ProgressResponse)
+        self.app.add_api_route("/sdapi/v1/interrogate", self.interrogateapi, methods=["POST"])
         self.app.add_api_route("/sdapi/v1/interrupt", self.interruptapi, methods=["POST"])
         self.app.add_api_route("/sdapi/v1/options", self.get_config, methods=["GET"], response_model=OptionsModel)
         self.app.add_api_route("/sdapi/v1/options", self.set_config, methods=["POST"])
@@ -214,6 +215,19 @@ class Api:
 
         return ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image)
 
+    def interrogateapi(self, interrogatereq: InterrogateRequest):
+        image_b64 = interrogatereq.image
+        if image_b64 is None:
+            raise HTTPException(status_code=404, detail="Image not found") 
+
+        img = self.__base64_to_image(image_b64)
+
+        # Override object param
+        with self.queue_lock:
+            processed = shared.interrogator.interrogate(img)
+        
+        return InterrogateResponse(caption=processed)
+
     def interruptapi(self):
         shared.state.interrupt()
 

modules/api/models.py

@@ -65,6 +65,7 @@ class PydanticModelGenerator:
 
         self._model_name = model_name
         self._class_data = merge_class_params(class_instance)
+
         self._model_def = [
             ModelDef(
                 field=underscore(k),
@@ -167,6 +168,12 @@ class ProgressResponse(BaseModel):
     state: dict = Field(title="State", description="The current state snapshot")
     current_image: str = Field(default=None, title="Current image", description="The current image in base64 format. opts.show_progress_every_n_steps is required for this to work.")
 
+class InterrogateRequest(BaseModel):
+    image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
+
+class InterrogateResponse(BaseModel):
+    caption: str = Field(default=None, title="Caption", description="The generated caption for the image.")
+
 fields = {}
 for key, value in opts.data.items():
     metadata = opts.data_labels.get(key)
@@ -231,3 +238,4 @@ class ArtistItem(BaseModel):
     name: str = Field(title="Name")
     score: float = Field(title="Score")
     category: str = Field(title="Category")
+

modules/ldsr_model_arch.py

@@ -101,8 +101,8 @@ class LDSR:
         down_sample_rate = target_scale / 4
         wd = width_og * down_sample_rate
         hd = height_og * down_sample_rate
-        width_downsampled_pre = int(wd)
-        height_downsampled_pre = int(hd)
+        width_downsampled_pre = int(np.ceil(wd))
+        height_downsampled_pre = int(np.ceil(hd))
 
         if down_sample_rate != 1:
             print(
@@ -110,7 +110,12 @@ class LDSR:
             im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS)
         else:
             print(f"Down sample rate is 1 from {target_scale} / 4 (Not downsampling)")
-        logs = self.run(model["model"], im_og, diffusion_steps, eta)
+        
+        # pad width and height to multiples of 64, pads with the edge values of image to avoid artifacts
+        pad_w, pad_h = np.max(((2, 2), np.ceil(np.array(im_og.size) / 64).astype(int)), axis=0) * 64 - im_og.size
+        im_padded = Image.fromarray(np.pad(np.array(im_og), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge'))
+        
+        logs = self.run(model["model"], im_padded, diffusion_steps, eta)
 
         sample = logs["sample"]
         sample = sample.detach().cpu()
@@ -120,6 +125,9 @@ class LDSR:
         sample = np.transpose(sample, (0, 2, 3, 1))
         a = Image.fromarray(sample[0])
 
+        # remove padding
+        a = a.crop((0, 0) + tuple(np.array(im_og.size) * 4))
+
         del model
         gc.collect()
         torch.cuda.empty_cache()

modules/localization.py

@@ -3,6 +3,7 @@ import os
 import sys
 import traceback
 
+
 localizations = {}
 
 
@@ -16,6 +17,11 @@ def list_localizations(dirname):
 
         localizations[fn] = os.path.join(dirname, file)
 
+    from modules import scripts
+    for file in scripts.list_scripts("localizations", ".json"):
+        fn, ext = os.path.splitext(file.filename)
+        localizations[fn] = file.path
+
 
 def localization_js(current_localization_name):
     fn = localizations.get(current_localization_name, None)

modules/safe.py

@@ -23,11 +23,18 @@ def encode(*args):
 
 
 class RestrictedUnpickler(pickle.Unpickler):
+    extra_handler = None
+
     def persistent_load(self, saved_id):
         assert saved_id[0] == 'storage'
         return TypedStorage()
 
     def find_class(self, module, name):
+        if self.extra_handler is not None:
+            res = self.extra_handler(module, name)
+            if res is not None:
+                return res
+
         if module == 'collections' and name == 'OrderedDict':
             return getattr(collections, name)
         if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter']:
@@ -52,7 +59,7 @@ class RestrictedUnpickler(pickle.Unpickler):
             return set
 
         # Forbid everything else.
-        raise pickle.UnpicklingError(f"global '{module}/{name}' is forbidden")
+        raise Exception(f"global '{module}/{name}' is forbidden")
 
 
 allowed_zip_names = ["archive/data.pkl", "archive/version"]
@@ -69,7 +76,7 @@ def check_zip_filenames(filename, names):
         raise Exception(f"bad file inside {filename}: {name}")
 
 
-def check_pt(filename):
+def check_pt(filename, extra_handler):
     try:
 
         # new pytorch format is a zip file
@@ -78,6 +85,7 @@ def check_pt(filename):
 
             with z.open('archive/data.pkl') as file:
                 unpickler = RestrictedUnpickler(file)
+                unpickler.extra_handler = extra_handler
                 unpickler.load()
 
     except zipfile.BadZipfile:
@@ -85,16 +93,42 @@ def check_pt(filename):
         # if it's not a zip file, it's an olf pytorch format, with five objects written to pickle
         with open(filename, "rb") as file:
             unpickler = RestrictedUnpickler(file)
+            unpickler.extra_handler = extra_handler
             for i in range(5):
                 unpickler.load()
 
 
 def load(filename, *args, **kwargs):
+    return load_with_extra(filename, *args, **kwargs)
+
+
+def load_with_extra(filename, extra_handler=None, *args, **kwargs):
+    """
+    this functon is intended to be used by extensions that want to load models with
+    some extra classes in them that the usual unpickler would find suspicious.
+
+    Use the extra_handler argument to specify a function that takes module and field name as text,
+    and returns that field's value:
+
+    ```python
+    def extra(module, name):
+        if module == 'collections' and name == 'OrderedDict':
+            return collections.OrderedDict
+
+        return None
+
+    safe.load_with_extra('model.pt', extra_handler=extra)
+    ```
+
+    The alternative to this is just to use safe.unsafe_torch_load('model.pt'), which as the name implies is
+    definitely unsafe.
+    """
+
     from modules import shared
 
     try:
         if not shared.cmd_opts.disable_safe_unpickle:
-            check_pt(filename)
+            check_pt(filename, extra_handler)
 
     except pickle.UnpicklingError:
         print(f"Error verifying pickled file from {filename}:", file=sys.stderr)

modules/script_callbacks.py

@@ -7,6 +7,7 @@ from typing import Optional
 from fastapi import FastAPI
 from gradio import Blocks
 
+
 def report_exception(c, job):
     print(f"Error executing callback {job} for {c.script}", file=sys.stderr)
     print(traceback.format_exc(), file=sys.stderr)
@@ -45,15 +46,21 @@ class CFGDenoiserParams:
         """Total number of sampling steps planned"""
 
 
+class UiTrainTabParams:
+    def __init__(self, txt2img_preview_params):
+        self.txt2img_preview_params = txt2img_preview_params
+
+
 ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])
 callback_map = dict(
     callbacks_app_started=[],
     callbacks_model_loaded=[],
     callbacks_ui_tabs=[],
+    callbacks_ui_train_tabs=[],
     callbacks_ui_settings=[],
     callbacks_before_image_saved=[],
     callbacks_image_saved=[],
-    callbacks_cfg_denoiser=[]
+    callbacks_cfg_denoiser=[],
 )
 
 
@@ -61,6 +68,7 @@ def clear_callbacks():
     for callback_list in callback_map.values():
         callback_list.clear()
 
+
 def app_started_callback(demo: Optional[Blocks], app: FastAPI):
     for c in callback_map['callbacks_app_started']:
         try:
@@ -89,6 +97,14 @@ def ui_tabs_callback():
     return res
 
 
+def ui_train_tabs_callback(params: UiTrainTabParams):
+    for c in callback_map['callbacks_ui_train_tabs']:
+        try:
+            c.callback(params)
+        except Exception:
+            report_exception(c, 'callbacks_ui_train_tabs')
+
+
 def ui_settings_callback():
     for c in callback_map['callbacks_ui_settings']:
         try:
@@ -169,6 +185,13 @@ def on_ui_tabs(callback):
     add_callback(callback_map['callbacks_ui_tabs'], callback)
 
 
+def on_ui_train_tabs(callback):
+    """register a function to be called when the UI is creating new tabs for the train tab.
+    Create your new tabs with gr.Tab.
+    """
+    add_callback(callback_map['callbacks_ui_train_tabs'], callback)
+
+
 def on_ui_settings(callback):
     """register a function to be called before UI settings are populated; add your settings
     by using shared.opts.add_option(shared.OptionInfo(...)) """

modules/scripts.py

@@ -3,7 +3,6 @@ import sys
 import traceback
 from collections import namedtuple
 
-import modules.ui as ui
 import gradio as gr
 
 from modules.processing import StableDiffusionProcessing

modules/shared.py

@@ -221,8 +221,6 @@ interrogator = modules.interrogate.InterrogateModels("interrogate")
 
 face_restorers = []
 
-localization.list_localizations(cmd_opts.localizations_dir)
-
 
 def realesrgan_models_names():
     import modules.realesrgan_model

modules/textual_inversion/preprocess.py

@@ -35,64 +35,62 @@ def preprocess(process_src, process_dst, process_width, process_height, preproce
             deepbooru.release_process()
 
 
+def listfiles(dirname):
+    return os.listdir(dirname)
 
-def preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_flip, process_split, process_caption, process_caption_deepbooru=False, split_threshold=0.5, overlap_ratio=0.2, process_focal_crop=False, process_focal_crop_face_weight=0.9, process_focal_crop_entropy_weight=0.3, process_focal_crop_edges_weight=0.5, process_focal_crop_debug=False):
-    width = process_width
-    height = process_height
-    src = os.path.abspath(process_src)
-    dst = os.path.abspath(process_dst)
-    split_threshold = max(0.0, min(1.0, split_threshold))
-    overlap_ratio = max(0.0, min(0.9, overlap_ratio))
-
-    assert src != dst, 'same directory specified as source and destination'
 
-    os.makedirs(dst, exist_ok=True)
+class PreprocessParams:
+    src = None
+    dstdir = None
+    subindex = 0
+    flip = False
+    process_caption = False
+    process_caption_deepbooru = False
+    preprocess_txt_action = None
 
-    files = os.listdir(src)
 
-    shared.state.textinfo = "Preprocessing..."
-    shared.state.job_count = len(files)
-
-    def save_pic_with_caption(image, index, existing_caption=None):
+def save_pic_with_caption(image, index, params: PreprocessParams, existing_caption=None):
     caption = ""
 
-        if process_caption:
+    if params.process_caption:
         caption += shared.interrogator.generate_caption(image)
 
-        if process_caption_deepbooru:
+    if params.process_caption_deepbooru:
         if len(caption) > 0:
             caption += ", "
         caption += deepbooru.get_tags_from_process(image)
 
-        filename_part = filename
+    filename_part = params.src
     filename_part = os.path.splitext(filename_part)[0]
     filename_part = os.path.basename(filename_part)
 
-        basename = f"{index:05}-{subindex[0]}-{filename_part}"
-        image.save(os.path.join(dst, f"{basename}.png"))
+    basename = f"{index:05}-{params.subindex}-{filename_part}"
+    image.save(os.path.join(params.dstdir, f"{basename}.png"))
 
-        if preprocess_txt_action == 'prepend' and existing_caption:
+    if params.preprocess_txt_action == 'prepend' and existing_caption:
         caption = existing_caption + ' ' + caption
-        elif preprocess_txt_action == 'append' and existing_caption:
+    elif params.preprocess_txt_action == 'append' and existing_caption:
         caption = caption + ' ' + existing_caption
-        elif preprocess_txt_action == 'copy' and existing_caption:
+    elif params.preprocess_txt_action == 'copy' and existing_caption:
         caption = existing_caption
 
     caption = caption.strip()
 
     if len(caption) > 0:
-            with open(os.path.join(dst, f"{basename}.txt"), "w", encoding="utf8") as file:
+        with open(os.path.join(params.dstdir, f"{basename}.txt"), "w", encoding="utf8") as file:
             file.write(caption)
 
-        subindex[0] += 1
+    params.subindex += 1
+
 
-    def save_pic(image, index, existing_caption=None):
-        save_pic_with_caption(image, index, existing_caption=existing_caption)
+def save_pic(image, index, params, existing_caption=None):
+    save_pic_with_caption(image, index, params, existing_caption=existing_caption)
 
-        if process_flip:
-            save_pic_with_caption(ImageOps.mirror(image), index, existing_caption=existing_caption)
+    if params.flip:
+        save_pic_with_caption(ImageOps.mirror(image), index, params, existing_caption=existing_caption)
 
-    def split_pic(image, inverse_xy):
+
+def split_pic(image, inverse_xy, width, height, overlap_ratio):
     if inverse_xy:
         from_w, from_h = image.height, image.width
         to_w, to_h = height, width
@@ -116,14 +114,40 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
         yield splitted
 
 
+def preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_flip, process_split, process_caption, process_caption_deepbooru=False, split_threshold=0.5, overlap_ratio=0.2, process_focal_crop=False, process_focal_crop_face_weight=0.9, process_focal_crop_entropy_weight=0.3, process_focal_crop_edges_weight=0.5, process_focal_crop_debug=False):
+    width = process_width
+    height = process_height
+    src = os.path.abspath(process_src)
+    dst = os.path.abspath(process_dst)
+    split_threshold = max(0.0, min(1.0, split_threshold))
+    overlap_ratio = max(0.0, min(0.9, overlap_ratio))
+
+    assert src != dst, 'same directory specified as source and destination'
+
+    os.makedirs(dst, exist_ok=True)
+
+    files = listfiles(src)
+
+    shared.state.textinfo = "Preprocessing..."
+    shared.state.job_count = len(files)
+
+    params = PreprocessParams()
+    params.dstdir = dst
+    params.flip = process_flip
+    params.process_caption = process_caption
+    params.process_caption_deepbooru = process_caption_deepbooru
+    params.preprocess_txt_action = preprocess_txt_action
+
     for index, imagefile in enumerate(tqdm.tqdm(files)):
-        subindex = [0]
+        params.subindex = 0
         filename = os.path.join(src, imagefile)
         try:
             img = Image.open(filename).convert("RGB")
         except Exception:
             continue
 
+        params.src = filename
+
         existing_caption = None
         existing_caption_filename = os.path.splitext(filename)[0] + '.txt'
         if os.path.exists(existing_caption_filename):
@@ -143,8 +167,8 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
         process_default_resize = True
 
         if process_split and ratio < 1.0 and ratio <= split_threshold:
-            for splitted in split_pic(img, inverse_xy):
-                save_pic(splitted, index, existing_caption=existing_caption)
+            for splitted in split_pic(img, inverse_xy, width, height, overlap_ratio):
+                save_pic(splitted, index, params, existing_caption=existing_caption)
             process_default_resize = False
 
         if process_focal_crop and img.height != img.width:
@@ -165,11 +189,11 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
                 dnn_model_path = dnn_model_path,
             )
             for focal in autocrop.crop_image(img, autocrop_settings):
-                save_pic(focal, index, existing_caption=existing_caption)
+                save_pic(focal, index, params, existing_caption=existing_caption)
             process_default_resize = False
 
         if process_default_resize:
             img = images.resize_image(1, img, width, height)
-            save_pic(img, index, existing_caption=existing_caption)
+            save_pic(img, index, params, existing_caption=existing_caption)
 
         shared.state.nextjob()

requirements.txt

@@ -4,7 +4,7 @@ fairscale==0.4.4
 fonts
 font-roboto
 gfpgan
-gradio==3.8
+gradio==3.9
 invisible-watermark
 numpy
 omegaconf