VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-09 06:45:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix sampling in training with mutiple gpus ref #989

Browse Source
This commit is contained in:
Kohya S
2023-12-15 22:35:54 +09:00
parent 7adcd9cd1a
commit 3efd90b2ad
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
library/train_util.py
View File

@@ -4603,7 +4603,7 @@ def line_to_prompt_dict(line: str) -> dict:
def sample_images_common( def sample_images_common(
pipe_class, pipe_class,
accelerator, accelerator: Accelerator,
args: argparse.Namespace, args: argparse.Namespace,
epoch, epoch,
steps, steps,
@@ -4640,6 +4640,13 @@ def sample_images_common(
org_vae_device = vae.device # CPUにいるはず org_vae_device = vae.device # CPUにいるはず
vae.to(device) vae.to(device)
# unwrap unet and text_encoder(s)
unet = accelerator.unwrap_model(unet)
if isinstance(text_encoder, (list, tuple)):
text_encoder = [accelerator.unwrap_model(te) for te in text_encoder]
else:
text_encoder = accelerator.unwrap_model(text_encoder)
# read prompts # read prompts
# with open(args.sample_prompts, "rt", encoding="utf-8") as f: # with open(args.sample_prompts, "rt", encoding="utf-8") as f: