Inference mode fix
This commit is contained in:
space-nuko
@@ -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
|
||||
|
||||
|
||||
+9
-8
@@ -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
|
||||
|
||||
+29
-28
@@ -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)
|
||||
|
||||
@@ -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,)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user