Skip to content

GitLab

"doc/vscode:/vscode.git/clone" did not exist on "85aeae712ab3ec1193c6c503ae06352892dcd9df"
Commit 149a4de3 authored by comfyanonymous's avatar comfyanonymous
Browse files

Fix potential issue if exception happens when patching model.

parent ef90e9c3
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
nodes.py
View file @ 149a4de3
...@@ -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,6 +323,7 @@ def common_ksampler(device, model, seed, steps, cfg, sampler_name, scheduler, po ...@@ -322,6 +323,7 @@ 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:
if real_model is not None:
real_model.cpu() 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