Commit 149a4de3 authored by comfyanonymous's avatar comfyanonymous
Browse files

Fix potential issue if exception happens when patching model.

parent ef90e9c3
...@@ -289,6 +289,7 @@ def common_ksampler(device, model, seed, steps, cfg, sampler_name, scheduler, po ...@@ -289,6 +289,7 @@ def common_ksampler(device, model, seed, steps, cfg, sampler_name, scheduler, po
else: else:
noise = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=torch.manual_seed(seed), device="cpu") noise = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=torch.manual_seed(seed), device="cpu")
real_model = None
try: try:
real_model = model.patch_model() real_model = model.patch_model()
real_model.to(device) real_model.to(device)
...@@ -322,7 +323,8 @@ def common_ksampler(device, model, seed, steps, cfg, sampler_name, scheduler, po ...@@ -322,7 +323,8 @@ def common_ksampler(device, model, seed, steps, cfg, sampler_name, scheduler, po
real_model.cpu() real_model.cpu()
model.unpatch_model() model.unpatch_model()
except Exception as e: except Exception as e:
real_model.cpu() if real_model is not None:
real_model.cpu()
model.unpatch_model() model.unpatch_model()
raise e raise e
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment