diff --git a/web/models/stable_diffusion/cache_objects.py b/web/models/stable_diffusion/cache_objects.py index abe5587415..e24356f158 100644 --- a/web/models/stable_diffusion/cache_objects.py +++ b/web/models/stable_diffusion/cache_objects.py @@ -16,6 +16,7 @@ from models.stable_diffusion.schedulers import ( SharkEulerDiscreteScheduler, ) +import gc model_config = { @@ -81,6 +82,9 @@ def __init__(self): self.version = None self.schedulers = None self.tokenizer = None + self.vae = None + self.clip = None + self.unet = None def set_models(self, device_key): if self.device != device_key or self.variant != args.variant: @@ -91,6 +95,12 @@ def set_models(self, device_key): args.max_length = 64 args.use_tuned = True set_init_device_flags() + del self.schedulers + del self.tokenizer + del self.vae + del self.unet + del self.clip + gc.collect() self.schedulers = get_schedulers(args.version) self.tokenizer = get_tokenizer(args.version) self.vae = get_vae()