Allow different models to estimate memory usage differently.

This commit is contained in:
comfyanonymous
2023-11-12 04:02:16 -05:00
parent 2c9dba8dc0
commit dd4ba68b6e
5 changed files with 19 additions and 26 deletions

View File

@@ -157,6 +157,16 @@ class BaseModel(torch.nn.Module):
def set_inpaint(self):
self.inpaint_model = True
def memory_required(self, input_shape):
area = input_shape[0] * input_shape[2] * input_shape[3]
if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention():
#TODO: this needs to be tweaked
return (area / 20) * (1024 * 1024)
else:
#TODO: this formula might be too aggressive since I tweaked the sub-quad and split algorithms to use less memory.
return (((area * 0.6) / 0.9) + 1024) * (1024 * 1024)
def unclip_adm(unclip_conditioning, device, noise_augmentor, noise_augment_merge=0.0):
adm_inputs = []
weights = []