Update pipeline.py

pipeline.py

@@ -16,7 +16,11 @@
 # modified from https://github.com/AUTOMATIC1111/stable-diffusion-webui
 # Here is the AGPL-3.0 license https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/LICENSE.txt
 
+import copy
 import inspect
+import os
+import os.path
+import shutil
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -24,6 +28,7 @@ import paddle
 import paddle.nn as nn
 import PIL
 import PIL.Image
+from huggingface_hub.file_download import _request_wrapper, hf_raise_for_status
 
 from paddlenlp.transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
 from ppdiffusers.models import AutoencoderKL, ControlNetModel, UNet2DConditionModel
@@ -35,7 +40,9 @@ from ppdiffusers.pipelines.stable_diffusion.safety_checker import (
 from ppdiffusers.schedulers import KarrasDiffusionSchedulers
 from ppdiffusers.utils import (
     PIL_INTERPOLATION,
+    PPDIFFUSERS_CACHE,
     logging,
+    ppdiffusers_url_download,
     randn_tensor,
     safetensors_load,
     smart_load,
@@ -43,6 +50,55 @@ from ppdiffusers.utils import (
 )
 
 
+def get_civitai_download_url(display_url, url_prefix="https://civitai.com"):
+    if "api/download" in display_url:
+        return display_url
+    import bs4
+    import requests
+
+    headers = {
+        "User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36 QIHU 360SE"
+    }
+    r = requests.get(display_url, headers=headers)
+    soup = bs4.BeautifulSoup(r.text, "lxml")
+    download_url = None
+    for a in soup.find_all("a", href=True):
+        if "Download" in str(a):
+            download_url = url_prefix + a["href"].split("?")[0]
+            break
+    return download_url
+
+
+def http_file_name(
+    url: str,
+    *,
+    proxies=None,
+    headers: Optional[Dict[str, str]] = None,
+    timeout=10.0,
+    max_retries=0,
+):
+    """
+    Get a remote file name.
+    """
+    headers = copy.deepcopy(headers) or {}
+    r = _request_wrapper(
+        method="GET",
+        url=url,
+        stream=True,
+        proxies=proxies,
+        headers=headers,
+        timeout=timeout,
+        max_retries=max_retries,
+    )
+    hf_raise_for_status(r)
+    displayed_name = url
+    content_disposition = r.headers.get("Content-Disposition")
+    if content_disposition is not None and "filename=" in content_disposition:
+        # Means file is on CDN
+        displayed_name = content_disposition.split("filename=")[-1]
+    return displayed_name
+
+
 @paddle.no_grad()
 def load_lora(
     pipeline,
@@ -168,6 +224,9 @@ class WebUIStableDiffusionControlNetPipeline(DiffusionPipeline):
     enable_emphasis = True
     comma_padding_backtrack = 20
 
+    LORA_DIR = os.path.join(PPDIFFUSERS_CACHE, "lora")
+    TI_DIR = os.path.join(PPDIFFUSERS_CACHE, "textual_inversion")
+
     def __init__(
         self,
         vae: AutoencoderKL,
@@ -232,7 +291,17 @@ class WebUIStableDiffusionControlNetPipeline(DiffusionPipeline):
         ]
         self.weights_has_changed = False
 
-    def add_ti_embedding_dir(self, embeddings_dir):
+        # register_state_dict_hook to fix text_encoder, when we save_pretrained text model.
+        def map_to(state_dict, *args, **kwargs):
+            if "text_model.token_embedding.wrapped.weight" in state_dict:
+                state_dict["text_model.token_embedding.weight"] = state_dict.pop(
+                    "text_model.token_embedding.wrapped.weight"
+                )
+            return state_dict
+
+        self.text_encoder.register_state_dict_hook(map_to)
+
+    def add_ti_embedding_dir(self, embeddings_dir=None):
         self.sj.embedding_db.add_embedding_dir(embeddings_dir)
         self.sj.embedding_db.load_textual_inversion_embeddings()
 
@@ -240,6 +309,30 @@ class WebUIStableDiffusionControlNetPipeline(DiffusionPipeline):
         self.sj.embedding_db.clear_embedding_dirs()
         self.sj.embedding_db.load_textual_inversion_embeddings(True)
 
+    def download_civitai_lora_file(self, url):
+        if os.path.isfile(url):
+            dst = os.path.join(self.LORA_DIR, os.path.basename(url))
+            shutil.copyfile(url, dst)
+            return dst
+
+        download_url = get_civitai_download_url(url) or url
+        file_path = ppdiffusers_url_download(
+            download_url, cache_dir=self.LORA_DIR, filename=http_file_name(download_url).strip('"')
+        )
+        return file_path
+
+    def download_civitai_ti_file(self, url):
+        if os.path.isfile(url):
+            dst = os.path.join(self.TI_DIR, os.path.basename(url))
+            shutil.copyfile(url, dst)
+            return dst
+
+        download_url = get_civitai_download_url(url) or url
+        file_path = ppdiffusers_url_download(
+            download_url, cache_dir=self.TI_DIR, filename=http_file_name(download_url).strip('"')
+        )
+        return file_path
+
     def change_scheduler(self, scheduler_type="ddim"):
         self.switch_scheduler(scheduler_type)
 
@@ -507,7 +600,6 @@ class WebUIStableDiffusionControlNetPipeline(DiffusionPipeline):
         cross_attention_kwargs: Optional[Dict[str, Any]] = None,
         clip_skip: int = 1,
         controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
-        lora_dir: str = "./loras",
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -571,8 +663,6 @@ class WebUIStableDiffusionControlNetPipeline(DiffusionPipeline):
                 The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
                 to the residual in the original unet. If multiple ControlNets are specified in init, you can set the
                 corresponding scale as a list.
-            lora_dir (`str`, *optional*):
-                Path to lora which we want to load.
         Examples:
 
         Returns:
@@ -582,6 +672,8 @@ class WebUIStableDiffusionControlNetPipeline(DiffusionPipeline):
             list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
             (nsfw) content, according to the `safety_checker`.
         """
+        self.add_ti_embedding_dir(self.TI_DIR)
+
         try:
             # 0. Default height and width to unet
             height, width = self._default_height_width(height, width, image)
@@ -613,19 +705,23 @@ class WebUIStableDiffusionControlNetPipeline(DiffusionPipeline):
 
             prompts, extra_network_data = parse_prompts([prompt])
 
-            if lora_dir is not None and os.path.exists(lora_dir):
-                lora_mapping = {p.stem: p.absolute() for p in Path(lora_dir).glob("*.safetensors")}
-                for params in extra_network_data["lora"]:
-                    assert len(params.items) > 0
-                    name = params.items[0]
-                    if name in lora_mapping:
-                        ratio = float(params.items[1]) if len(params.items) > 1 else 1.0
-                        lora_state_dict = smart_load(lora_mapping[name], map_location=paddle.get_device())
-                        self.weights_has_changed = True
-                        load_lora(self, state_dict=lora_state_dict, ratio=ratio)
-                        del lora_state_dict
-                    else:
-                        print(f"We can't find lora weight: {name}! Please make sure that exists!")
+            if self.LORA_DIR is not None:
+                if os.path.exists(self.LORA_DIR):
+                    lora_mapping = {p.stem: p.absolute() for p in Path(self.LORA_DIR).glob("*.safetensors")}
+                    for params in extra_network_data["lora"]:
+                        assert len(params.items) > 0
+                        name = params.items[0]
+                        if name in lora_mapping:
+                            ratio = float(params.items[1]) if len(params.items) > 1 else 1.0
+                            lora_state_dict = smart_load(lora_mapping[name], map_location=paddle.get_device())
+                            self.weights_has_changed = True
+                            load_lora(self, state_dict=lora_state_dict, ratio=ratio)
+                            del lora_state_dict
+                        else:
+                            print(f"We can't find lora weight: {name}! Please make sure that exists!")
+                else:
+                    if len(extra_network_data["lora"]) > 0:
+                        print(f"{self.LORA_DIR} not exists, so we cant load loras!")
 
             self.sj.clip.CLIP_stop_at_last_layers = clip_skip
             # 3. Encode input prompt
@@ -1808,7 +1904,7 @@ class EmbeddingDatabase:
         self.previously_displayed_embeddings = ()
 
     def add_embedding_dir(self, path):
-        if path is not None:
+        if path is not None and path not in self.embedding_dirs:
             self.embedding_dirs[path] = DirWithTextualInversionEmbeddings(path)
 
     def clear_embedding_dirs(self):