reduce peak VRAM in sample gen

2026-04-08 14:34:23 +00:00 · 2024-02-04 17:31:01 +09:00
parent 2f9a344297
commit e793d7780d
1 changed files with 4 additions and 0 deletions
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -4820,6 +4820,10 @@ def sample_image_inference(accelerator: Accelerator, args: argparse.Namespace, p
            controlnet=controlnet,
            controlnet_image=controlnet_image,
        )
+    
+    with torch.cuda.device(torch.cuda.current_device()):
+        torch.cuda.empty_cache()
+        
    image = pipeline.latents_to_image(latents)[0]

    # adding accelerator.wait_for_everyone() here should sync up and ensure that sample images are saved in the same order as the original prompt list