Merge branch 'master' into patch-1

7 files changed, 74 insertions, 16 deletions

diff --git a/modules/api/api.py b/modules/api/api.py
index 688469ad..596a6616 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -15,6 +15,9 @@ from modules.sd_models import checkpoints_list
 from modules.realesrgan_model import get_realesrgan_models
 from typing import List
 
+if shared.cmd_opts.deepdanbooru:
+    from modules.deepbooru import get_deepbooru_tags
+
 def upscaler_to_index(name: str):
     try:
         return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())
@@ -220,11 +223,20 @@ class Api:
         if image_b64 is None:
             raise HTTPException(status_code=404, detail="Image not found") 
 
-        img = self.__base64_to_image(image_b64)
+        img = decode_base64_to_image(image_b64)
+        img = img.convert('RGB')
 
         # Override object param
         with self.queue_lock:
-            processed = shared.interrogator.interrogate(img)
+            if interrogatereq.model == "clip":
+                processed = shared.interrogator.interrogate(img)
+            elif interrogatereq.model == "deepdanbooru":
+                if shared.cmd_opts.deepdanbooru:
+                    processed = get_deepbooru_tags(img)
+                else:
+                    raise HTTPException(status_code=404, detail="Model not found. Add --deepdanbooru when launching for using the model.")
+            else:
+                raise HTTPException(status_code=404, detail="Model not found")
         
         return InterrogateResponse(caption=processed)
 
diff --git a/modules/api/models.py b/modules/api/models.py
index 34dbfa16..f9cd929e 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -170,6 +170,7 @@ class ProgressResponse(BaseModel):
 
 class InterrogateRequest(BaseModel):
     image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
+    model: str = Field(default="clip", title="Model", description="The interrogate model used.")
 
 class InterrogateResponse(BaseModel):
     caption: str = Field(default=None, title="Caption", description="The generated caption for the image.")
diff --git a/modules/ngrok.py b/modules/ngrok.py
index 25c53af8..64c9a3c2 100644
--- a/modules/ngrok.py
+++ b/modules/ngrok.py
@@ -1,14 +1,23 @@
 from pyngrok import ngrok, conf, exception
 
-
 def connect(token, port, region):
+    account = None
     if token == None:
         token = 'None'
+    else:
+        if ':' in token:
+            # token = authtoken:username:password
+            account = token.split(':')[1] + ':' + token.split(':')[-1]
+            token = token.split(':')[0]
+
     config = conf.PyngrokConfig(
         auth_token=token, region=region
     )
     try:
-        public_url = ngrok.connect(port, pyngrok_config=config, bind_tls=True).public_url
+        if account == None:
+            public_url = ngrok.connect(port, pyngrok_config=config, bind_tls=True).public_url
+        else:
+            public_url = ngrok.connect(port, pyngrok_config=config, bind_tls=True, auth=account).public_url
     except exception.PyngrokNgrokError:
         print(f'Invalid ngrok authtoken, ngrok connection aborted.\n'
               f'Your token: {token}, get the right one on https://dashboard.ngrok.com/get-started/your-authtoken')
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 34c57bfa..80addf03 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -163,13 +163,21 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
     checkpoint_file = checkpoint_info.filename

     sd_model_hash = checkpoint_info.hash

 

-    if shared.opts.sd_checkpoint_cache > 0 and hasattr(model, "sd_checkpoint_info"):

+    cache_enabled = shared.opts.sd_checkpoint_cache > 0

+

+    if cache_enabled:

         sd_vae.restore_base_vae(model)

-        checkpoints_loaded[model.sd_checkpoint_info] = model.state_dict().copy()

 

     vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)

 

-    if checkpoint_info not in checkpoints_loaded:

+    if cache_enabled and checkpoint_info in checkpoints_loaded:

+        # use checkpoint cache

+        vae_name = sd_vae.get_filename(vae_file) if vae_file else None

+        vae_message = f" with {vae_name} VAE" if vae_name else ""

+        print(f"Loading weights [{sd_model_hash}]{vae_message} from cache")

+        model.load_state_dict(checkpoints_loaded[checkpoint_info])

+    else:

+        # load from file

         print(f"Loading weights [{sd_model_hash}] from {checkpoint_file}")

 

         pl_sd = torch.load(checkpoint_file, map_location=shared.weight_load_location)

@@ -180,6 +188,10 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
         del pl_sd

         model.load_state_dict(sd, strict=False)

         del sd

+        

+        if cache_enabled:

+            # cache newly loaded model

+            checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

 

         if shared.cmd_opts.opt_channelslast:

             model.to(memory_format=torch.channels_last)

@@ -199,14 +211,9 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
 

         model.first_stage_model.to(devices.dtype_vae)

 

-    else:

-        vae_name = sd_vae.get_filename(vae_file) if vae_file else None

-        vae_message = f" with {vae_name} VAE" if vae_name else ""

-        print(f"Loading weights [{sd_model_hash}]{vae_message} from cache")

-        model.load_state_dict(checkpoints_loaded[checkpoint_info])

-

-    if shared.opts.sd_checkpoint_cache > 0:

-        while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:

+    # clean up cache if limit is reached

+    if cache_enabled:

+        while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache + 1: # we need to count the current model

             checkpoints_loaded.popitem(last=False)  # LRU

 

     model.sd_model_hash = sd_model_hash

diff --git a/modules/shared.py b/modules/shared.py
index e8bacd3c..caabf078 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -319,6 +319,8 @@ options_templates.update(options_section(('system', "System"), {
 

 options_templates.update(options_section(('training', "Training"), {

     "unload_models_when_training": OptionInfo(False, "Move VAE and CLIP to RAM when training if possible. Saves VRAM."),

+    "shuffle_tags": OptionInfo(False, "Shuffleing tags by ',' when create texts."),

+    "tag_drop_out": OptionInfo(0, "Dropout tags when create texts", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.1}),

     "save_optimizer_state": OptionInfo(False, "Saves Optimizer state as separate *.optim file. Training can be resumed with HN itself and matching optim file."),

     "dataset_filename_word_regex": OptionInfo("", "Filename word regex"),

     "dataset_filename_join_string": OptionInfo(" ", "Filename join string"),

diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index ad726577..eb75c376 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -98,7 +98,12 @@ class PersonalizedBase(Dataset):
     def create_text(self, filename_text):

         text = random.choice(self.lines)

         text = text.replace("[name]", self.placeholder_token)

-        text = text.replace("[filewords]", filename_text)

+        tags = filename_text.split(',')

+        if shared.opts.tag_drop_out != 0:

+            tags = [t for t in tags if random.random() > shared.opts.tag_drop_out]

+        if shared.opts.shuffle_tags:

+            random.shuffle(tags)

+        text = text.replace("[filewords]", ','.join(tags))

         return text

 

     def __len__(self):

diff --git a/modules/ui.py b/modules/ui.py
index 7ea1177f..5dce7f3b 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -566,6 +566,19 @@ def apply_setting(key, value):
     return value

 

 

+def update_generation_info(args):

+    generation_info, html_info, img_index = args

+    try:

+        generation_info = json.loads(generation_info)

+        if img_index < 0 or img_index >= len(generation_info["infotexts"]):

+            return html_info

+        return plaintext_to_html(generation_info["infotexts"][img_index])

+    except Exception:

+        pass

+    # if the json parse or anything else fails, just return the old html_info

+    return html_info

+

+

 def create_refresh_button(refresh_component, refresh_method, refreshed_args, elem_id):

     def refresh():

         refresh_method()

@@ -638,6 +651,15 @@ Requested path was: {f}
                     with gr.Group():

                         html_info = gr.HTML()

                         generation_info = gr.Textbox(visible=False)

+                        if tabname == 'txt2img' or tabname == 'img2img':

+                            generation_info_button = gr.Button(visible=False, elem_id=f"{tabname}_generation_info_button")

+                            generation_info_button.click(

+                                fn=update_generation_info,

+                                _js="(x, y) => [x, y, selected_gallery_index()]",

+                                inputs=[generation_info, html_info],

+                                outputs=[html_info],

+                                preprocess=False

+                            )

 

                         save.click(

                             fn=wrap_gradio_call(save_files),