diff --git a/README.md b/README.md index d236df0..4803414 100644 --- a/README.md +++ b/README.md @@ -9,19 +9,7 @@ This repo holds a modularized version of [Disco Diffusion](https://github.com/al 1. Clone this repository into your `ComfyUI/custom_nodes` folder with `git clone --recursive https://github.com/space-nuko/ComfyUI-Disco-Diffusion`. 2. Activate the virtualenv you use with ComfyUI. 3. Navigate to the cloned folder and run `pip install -r requirements.txt`. - -**NOTE:** You also have to make sure that inference mode is turned off within ComfyUI or you'll get errors. - -In `ComfyUI/execution.py`, change: - -``` python -with torch.inference_mode(): -``` - -to: -``` python -with torch.inference_mode(False): -``` +4. Start ComfyUI. ## Nodes included diff --git a/diffuse.py b/diffuse.py index 4e1a7c7..1278d4c 100644 --- a/diffuse.py +++ b/diffuse.py @@ -241,15 +241,16 @@ def diffuse(model, diffusion, clip_model, clip_vision, args: DiscoDiffusionSetti results = [] - gc.collect() - torch.cuda.empty_cache() - try: - results = do_run(diffusion, model, clip_model, clip_vision, args, batchNum) - except KeyboardInterrupt: - pass - finally: - print('Seed used:', seed) + with torch.inference_mode(False): gc.collect() torch.cuda.empty_cache() + try: + results = do_run(diffusion, model, clip_model, clip_vision, args, batchNum) + except KeyboardInterrupt: + pass + finally: + print('Seed used:', seed) + gc.collect() + torch.cuda.empty_cache() return results diff --git a/model_settings.py b/model_settings.py index 63629f6..c128592 100644 --- a/model_settings.py +++ b/model_settings.py @@ -227,33 +227,34 @@ class ModelSettings: device = comfy.model_management.get_torch_device() - if self.use_secondary_model: - self.secondary_model = SecondaryDiffusionImageNet2() - self.secondary_model.load_state_dict(torch.load(f'{self.model_path}/secondary_model_imagenet_2.pth', map_location='cpu')) - self.secondary_model.eval().requires_grad_(False).to(device) + with torch.inference_mode(False): + if self.use_secondary_model: + self.secondary_model = SecondaryDiffusionImageNet2() + self.secondary_model.load_state_dict(torch.load(f'{self.model_path}/secondary_model_imagenet_2.pth', map_location='cpu')) + self.secondary_model.eval().requires_grad_(False).to(device) - self.clip_models = [] - #if self.ViTB32: clip_models.append(clip.load('ViT-B/32', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.ViTB16: clip_models.append(clip.load('ViT-B/16', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.ViTL14: clip_models.append(clip.load('ViT-L/14', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.ViTL14_336px: clip_models.append(clip.load('ViT-L/14@336px', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.RN50: clip_models.append(clip.load('RN50', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.RN50x4: clip_models.append(clip.load('RN50x4', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.RN50x16: clip_models.append(clip.load('RN50x16', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.RN50x64: clip_models.append(clip.load('RN50x64', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.RN101: clip_models.append(clip.load('RN101', jit=False)[0].eval().requires_grad_(False).to(device)) - #if self.ViTB32_laion2b_e16: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion2b_e16').eval().requires_grad_(False).to(device)) - #if self.ViTB32_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion400m_e31').eval().requires_grad_(False).to(device)) - #if self.ViTB32_laion400m_32: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion400m_e32').eval().requires_grad_(False).to(device)) - #if self.ViTB32quickgelu_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-32-quickgelu', pretrained='laion400m_e31').eval().requires_grad_(False).to(device)) - #if self.ViTB32quickgelu_laion400m_e32: clip_models.append(open_clip.create_model('ViT-B-32-quickgelu', pretrained='laion400m_e32').eval().requires_grad_(False).to(device)) - #if self.ViTB16_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-16', pretrained='laion400m_e31').eval().requires_grad_(False).to(device)) - #if self.ViTB16_laion400m_e32: clip_models.append(open_clip.create_model('ViT-B-16', pretrained='laion400m_e32').eval().requires_grad_(False).to(device)) - #if self.RN50_yffcc15m: clip_models.append(open_clip.create_model('RN50', pretrained='yfcc15m').eval().requires_grad_(False).to(device)) - #if self.RN50_cc12m: clip_models.append(open_clip.create_model('RN50', pretrained='cc12m').eval().requires_grad_(False).to(device)) - #if self.RN50_quickgelu_yfcc15m: clip_models.append(open_clip.create_model('RN50-quickgelu', pretrained='yfcc15m').eval().requires_grad_(False).to(device)) - #if self.RN50_quickgelu_cc12m: clip_models.append(open_clip.create_model('RN50-quickgelu', pretrained='cc12m').eval().requires_grad_(False).to(device)) - #if self.RN101_yfcc15m: clip_models.append(open_clip.create_model('RN101', pretrained='yfcc15m').eval().requires_grad_(False).to(device)) - #if self.RN101_quickgelu_yfcc15m: clip_models.append(open_clip.create_model('RN101-quickgelu', pretrained='yfcc15m').eval().requires_grad_(False).to(device)) + self.clip_models = [] + #if self.ViTB32: clip_models.append(clip.load('ViT-B/32', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.ViTB16: clip_models.append(clip.load('ViT-B/16', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.ViTL14: clip_models.append(clip.load('ViT-L/14', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.ViTL14_336px: clip_models.append(clip.load('ViT-L/14@336px', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.RN50: clip_models.append(clip.load('RN50', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.RN50x4: clip_models.append(clip.load('RN50x4', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.RN50x16: clip_models.append(clip.load('RN50x16', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.RN50x64: clip_models.append(clip.load('RN50x64', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.RN101: clip_models.append(clip.load('RN101', jit=False)[0].eval().requires_grad_(False).to(device)) + #if self.ViTB32_laion2b_e16: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion2b_e16').eval().requires_grad_(False).to(device)) + #if self.ViTB32_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion400m_e31').eval().requires_grad_(False).to(device)) + #if self.ViTB32_laion400m_32: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion400m_e32').eval().requires_grad_(False).to(device)) + #if self.ViTB32quickgelu_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-32-quickgelu', pretrained='laion400m_e31').eval().requires_grad_(False).to(device)) + #if self.ViTB32quickgelu_laion400m_e32: clip_models.append(open_clip.create_model('ViT-B-32-quickgelu', pretrained='laion400m_e32').eval().requires_grad_(False).to(device)) + #if self.ViTB16_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-16', pretrained='laion400m_e31').eval().requires_grad_(False).to(device)) + #if self.ViTB16_laion400m_e32: clip_models.append(open_clip.create_model('ViT-B-16', pretrained='laion400m_e32').eval().requires_grad_(False).to(device)) + #if self.RN50_yffcc15m: clip_models.append(open_clip.create_model('RN50', pretrained='yfcc15m').eval().requires_grad_(False).to(device)) + #if self.RN50_cc12m: clip_models.append(open_clip.create_model('RN50', pretrained='cc12m').eval().requires_grad_(False).to(device)) + #if self.RN50_quickgelu_yfcc15m: clip_models.append(open_clip.create_model('RN50-quickgelu', pretrained='yfcc15m').eval().requires_grad_(False).to(device)) + #if self.RN50_quickgelu_cc12m: clip_models.append(open_clip.create_model('RN50-quickgelu', pretrained='cc12m').eval().requires_grad_(False).to(device)) + #if self.RN101_yfcc15m: clip_models.append(open_clip.create_model('RN101', pretrained='yfcc15m').eval().requires_grad_(False).to(device)) + #if self.RN101_quickgelu_yfcc15m: clip_models.append(open_clip.create_model('RN101-quickgelu', pretrained='yfcc15m').eval().requires_grad_(False).to(device)) - self.lpips_model = lpips.LPIPS(net='vgg').to(device) + self.lpips_model = lpips.LPIPS(net='vgg').to(device) diff --git a/nodes.py b/nodes.py index c38e585..9bad565 100644 --- a/nodes.py +++ b/nodes.py @@ -47,13 +47,14 @@ class OpenAICLIPLoader: def load(self, model_name): device = comfy.model_management.get_torch_device() - if model_name in OPENAI_CLIP_MODELS: - clip_model = openai_clip.load(model_name, jit=False)[0] - else: - name, pretrained = model_name.split("_", 1) - clip_model = open_clip.create_model(name, pretrained=pretrained) + with torch.inference_mode(False): + if model_name in OPENAI_CLIP_MODELS: + clip_model = openai_clip.load(model_name, jit=False)[0] + else: + name, pretrained = model_name.split("_", 1) + clip_model = open_clip.create_model(name, pretrained=pretrained) - clip_model.eval().requires_grad_(False).to(device) + clip_model.eval().requires_grad_(False).to(device) return (clip_model, clip_model,)