feat: refactor latent cache format

kohya-ss · Nov 15, 2024 · bdac55e · bdac55e
1 parent 0047bb1
commit bdac55e
Show file tree

Hide file tree

Showing 17 changed files with 360 additions and 637 deletions.
diff --git a/fine_tune.py b/fine_tune.py
@@ -177,7 +177,7 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
         vae.requires_grad_(False)
         vae.eval()
 
-        train_dataset_group.new_cache_latents(vae, accelerator)
+        train_dataset_group.new_cache_latents(vae, accelerator, args.force_cache_precision)
 
         vae.to("cpu")
         clean_memory_on_device(accelerator.device)

diff --git a/finetune/prepare_buckets_latents.py b/finetune/prepare_buckets_latents.py
@@ -180,7 +180,7 @@ def process_batch(is_last):
 
         # バッチへ追加
         image_info = train_util.ImageInfo(image_key, 1, "", False, image_path)
-        image_info.latents_npz = npz_file_name
+        image_info.latents_cache_path = npz_file_name
         image_info.bucket_reso = reso
         image_info.resized_size = resized_size
         image_info.image = image

diff --git a/flux_train.py b/flux_train.py
@@ -198,7 +198,7 @@ def train(args):
         ae.requires_grad_(False)
         ae.eval()
 
-        train_dataset_group.new_cache_latents(ae, accelerator)
+        train_dataset_group.new_cache_latents(ae, accelerator, args.force_cache_precision)
 
         ae.to("cpu")  # if no sampling, vae can be deleted
         clean_memory_on_device(accelerator.device)