Generate sample images without having CUDA (such as on Macs)

This commit is contained in:
Millie
2024-07-18 11:52:58 -07:00
parent 25f961bc77
commit 2e67978ee2

View File

@@ -5229,7 +5229,7 @@ def sample_images_common(
clean_memory_on_device(accelerator.device) clean_memory_on_device(accelerator.device)
torch.set_rng_state(rng_state) torch.set_rng_state(rng_state)
if cuda_rng_state is not None: if torch.cuda.is_available() and cuda_rng_state is not None:
torch.cuda.set_rng_state(cuda_rng_state) torch.cuda.set_rng_state(cuda_rng_state)
vae.to(org_vae_device) vae.to(org_vae_device)
@@ -5263,11 +5263,13 @@ def sample_image_inference(
if seed is not None: if seed is not None:
torch.manual_seed(seed) torch.manual_seed(seed)
torch.cuda.manual_seed(seed) if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
else: else:
# True random sample image generation # True random sample image generation
torch.seed() torch.seed()
torch.cuda.seed() if torch.cuda.is_available():
torch.cuda.seed()
scheduler = get_my_scheduler( scheduler = get_my_scheduler(
sample_sampler=sampler_name, sample_sampler=sampler_name,
@@ -5302,8 +5304,9 @@ def sample_image_inference(
controlnet_image=controlnet_image, controlnet_image=controlnet_image,
) )
with torch.cuda.device(torch.cuda.current_device()): if torch.cuda.is_available():
torch.cuda.empty_cache() with torch.cuda.device(torch.cuda.current_device()):
torch.cuda.empty_cache()
image = pipeline.latents_to_image(latents)[0] image = pipeline.latents_to_image(latents)[0]