diff options
| -rw-r--r-- | modules/devices.py | 3 | ||||
| -rw-r--r-- | modules/processing.py | 21 |
2 files changed, 9 insertions, 15 deletions
diff --git a/modules/devices.py b/modules/devices.py index 6db4e57c..0158b11f 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -1,7 +1,6 @@ import contextlib import torch -import gc from modules import errors @@ -20,8 +19,8 @@ def get_optimal_device(): return cpu + def torch_gc(): - gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() diff --git a/modules/processing.py b/modules/processing.py index e7f9c85e..f666ba81 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -345,8 +345,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed: if state.job_count == -1:
state.job_count = p.n_iter
- for n in range(p.n_iter):
- with torch.no_grad(), precision_scope("cuda"), ema_scope():
+ for n in range(p.n_iter):
if state.interrupted:
break
@@ -395,22 +394,19 @@ def process_images(p: StableDiffusionProcessing) -> Processed: import modules.safety as safety
x_samples_ddim = modules.safety.censor_batch(x_samples_ddim)
- for i, x_sample in enumerate(x_samples_ddim):
- with torch.no_grad(), precision_scope("cuda"), ema_scope():
+ for i, x_sample in enumerate(x_samples_ddim):
x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
x_sample = x_sample.astype(np.uint8)
- if p.restore_faces:
- with torch.no_grad(), precision_scope("cuda"), ema_scope():
+ if p.restore_faces:
if opts.save and not p.do_not_save_samples and opts.save_images_before_face_restoration:
images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-face-restoration")
- x_sample = modules.face_restoration.restore_faces(x_sample)
devices.torch_gc()
- devices.torch_gc()
+ x_sample = modules.face_restoration.restore_faces(x_sample)
+ devices.torch_gc()
- with torch.no_grad(), precision_scope("cuda"), ema_scope():
image = Image.fromarray(x_sample)
if p.color_corrections is not None and i < len(p.color_corrections):
@@ -438,13 +434,12 @@ def process_images(p: StableDiffusionProcessing) -> Processed: infotexts.append(infotext(n, i))
output_images.append(image)
- del x_samples_ddim
+ del x_samples_ddim
- devices.torch_gc()
+ devices.torch_gc()
- state.nextjob()
+ state.nextjob()
- with torch.no_grad(), precision_scope("cuda"), ema_scope():
p.color_corrections = None
index_of_first_image = 0
|