migrate to new HF diffusers cache location

invokeai/backend/model_management/model_manager.py

@@ -43,13 +43,11 @@ class SDLegacyType(Enum):
     V2 = 3
     UNKNOWN = 99
 
-
 DEFAULT_MAX_MODELS = 2
 VAE_TO_REPO_ID = {  # hack, see note in convert_and_import()
     "vae-ft-mse-840000-ema-pruned": "stabilityai/sd-vae-ft-mse",
 }
 
-
 class ModelManager(object):
     def __init__(
         self,
@@ -369,7 +367,7 @@ class ModelManager(object):
             if vae := self._load_vae(mconfig["vae"]):
                 pipeline_args.update(vae=vae)
         if not isinstance(name_or_path, Path):
-            pipeline_args.update(cache_dir=global_cache_dir("diffusers"))
+            pipeline_args.update(cache_dir=global_cache_dir("hub"))
         if using_fp16:
             pipeline_args.update(torch_dtype=torch.float16)
             fp_args_list = [{"revision": "fp16"}, {}]
@@ -916,25 +914,30 @@ class ModelManager(object):
         to the 2.3.0 "diffusers" version. This should be a one-time operation, called at
         script startup time.
         """
-        # Three transformer models to check: bert, clip and safety checker
+        # Three transformer models to check: bert, clip and safety checker, and
+        # the diffusers as well
+        models_dir = Path(Globals.root, "models")
         legacy_locations = [
             Path(
+                models_dir,
                 "CompVis/stable-diffusion-safety-checker/models--CompVis--stable-diffusion-safety-checker"
             ),
-            Path("bert-base-uncased/models--bert-base-uncased"),
+            Path(models_dir, "bert-base-uncased/models--bert-base-uncased"),
             Path(
+                models_dir,
                 "openai/clip-vit-large-patch14/models--openai--clip-vit-large-patch14"
             ),
         ]
-        models_dir = Path(Globals.root, "models")
+        legacy_locations.extend(list(Path(models_dir,"diffusers").glob('*')))
+        
         legacy_layout = False
         for model in legacy_locations:
-            legacy_layout = legacy_layout or Path(models_dir, model).exists()
+            legacy_layout = legacy_layout or model.exists()
         if not legacy_layout:
             return
 
         print(
-            "** Legacy version <= 2.2.5 model directory layout detected. Reorganizing."
+            "** Old model directory layout (< v3.0) detected. Reorganizing."
         )
         print("** This is a quick one-time operation.")
 
@@ -948,6 +951,8 @@ class ModelManager(object):
         for model in legacy_locations:
             source = models_dir / model
             dest = hub / model.stem
+            if dest.exists() and not source.exists():
+                continue
             print(f"** {source} => {dest}")
             if source.exists():
                 if dest.exists():
@@ -955,26 +960,6 @@ class ModelManager(object):
                 else:
                     move(source, dest)
 
-        # anything else gets moved into the diffusers directory
-        if cls._is_huggingface_hub_directory_present():
-            diffusers = global_cache_dir("diffusers")
-        else:
-            diffusers = models_dir / "diffusers"
-
-        os.makedirs(diffusers, exist_ok=True)
-        for root, dirs, _ in os.walk(models_dir, topdown=False):
-            for dir in dirs:
-                full_path = Path(root, dir)
-                if full_path.is_relative_to(hub) or full_path.is_relative_to(diffusers):
-                    continue
-                if Path(dir).match("models--*--*"):
-                    dest = diffusers / dir
-                    print(f"** {full_path} => {dest}")
-                    if dest.exists():
-                        rmtree(full_path)
-                    else:
-                        move(full_path, dest)
-
         # now clean up by removing any empty directories
         empty = [
             root
@@ -1072,7 +1057,7 @@ class ModelManager(object):
             path = name_or_path
         else:
             owner, repo = name_or_path.split("/")
-            path = Path(global_cache_dir("diffusers") / f"models--{owner}--{repo}")
+            path = Path(global_cache_dir("hub") / f"models--{owner}--{repo}")
         if not path.exists():
             return None
         hashpath = path / "checksum.sha256"
@@ -1133,7 +1118,7 @@ class ModelManager(object):
         using_fp16 = self.precision == "float16"
 
         vae_args.update(
-            cache_dir=global_cache_dir("diffusers"),
+            cache_dir=global_cache_dir("hub"),
             local_files_only=not Globals.internet_available,
         )
 
@@ -1172,7 +1157,7 @@ class ModelManager(object):
 
     @staticmethod
     def _delete_model_from_cache(repo_id):
-        cache_info = scan_cache_dir(global_cache_dir("diffusers"))
+        cache_info = scan_cache_dir(global_cache_dir("hub"))
 
         # I'm sure there is a way to do this with comprehensions
         # but the code quickly became incomprehensible!