Quash memory leak when compel invocation called ()

This commit prevents each image generation from leaking ~160 MB of VRAM.
Thanks to @damian0815 and @StAlKeR7779 for helping to sort this out.
This commit is contained in:
blessedcoolant 2023-07-04 06:33:56 +12:00 committed by GitHub
commit b4b760d9e9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -2,6 +2,7 @@ from typing import Literal, Optional, Union
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from contextlib import ExitStack from contextlib import ExitStack
import re import re
import torch
from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig
from .model import ClipField from .model import ClipField
@ -56,6 +57,7 @@ class CompelInvocation(BaseInvocation):
}, },
} }
@torch.no_grad()
def invoke(self, context: InvocationContext) -> CompelOutput: def invoke(self, context: InvocationContext) -> CompelOutput:
tokenizer_info = context.services.model_manager.get_model( tokenizer_info = context.services.model_manager.get_model(