Skip to content

Commit 992712c

Browse files
authored
Merge pull request kohya-ss#1433 from millie-v/sample-image-without-cuda
Generate sample images without having CUDA (such as on Macs)
2 parents 6b2bb31 + 4165982 commit 992712c

File tree

1 file changed

+8
-5
lines changed

1 file changed

+8
-5
lines changed

library/train_util.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5404,7 +5404,7 @@ def sample_images_common(
54045404
clean_memory_on_device(accelerator.device)
54055405

54065406
torch.set_rng_state(rng_state)
5407-
if cuda_rng_state is not None:
5407+
if torch.cuda.is_available() and cuda_rng_state is not None:
54085408
torch.cuda.set_rng_state(cuda_rng_state)
54095409
vae.to(org_vae_device)
54105410

@@ -5438,11 +5438,13 @@ def sample_image_inference(
54385438

54395439
if seed is not None:
54405440
torch.manual_seed(seed)
5441-
torch.cuda.manual_seed(seed)
5441+
if torch.cuda.is_available():
5442+
torch.cuda.manual_seed(seed)
54425443
else:
54435444
# True random sample image generation
54445445
torch.seed()
5445-
torch.cuda.seed()
5446+
if torch.cuda.is_available():
5447+
torch.cuda.seed()
54465448

54475449
scheduler = get_my_scheduler(
54485450
sample_sampler=sampler_name,
@@ -5477,8 +5479,9 @@ def sample_image_inference(
54775479
controlnet_image=controlnet_image,
54785480
)
54795481

5480-
with torch.cuda.device(torch.cuda.current_device()):
5481-
torch.cuda.empty_cache()
5482+
if torch.cuda.is_available():
5483+
with torch.cuda.device(torch.cuda.current_device()):
5484+
torch.cuda.empty_cache()
54825485

54835486
image = pipeline.latents_to_image(latents)[0]
54845487

0 commit comments

Comments
 (0)