attempt to address memory issues when loading ckpt models (#2128)

- A couple of users have reported that switching back and forth
  between ckpt models is causing a "GPU out of memory" crash.
  Traceback suggests there is actually a CPU RAM issue.

- This speculative test simply performs a round of garbage collection
  before the point where the crash occurs.
This commit is contained in:
Lincoln Stein 2022-12-29 09:00:50 -05:00 committed by GitHub
parent 4cc60669c1
commit 9d103ef030
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -242,6 +242,9 @@ class ModelCache(object):
# merged models from auto11 merge board are flat for some reason
if 'state_dict' in sd:
sd = sd['state_dict']
print(f' | Forcing garbage collection prior to loading new model')
gc.collect()
model = instantiate_from_config(omega_config.model)
model.load_state_dict(sd, strict=False)