Update for 1.7.0
Update of the TRT Patch for the 1.7.0 version of the original file
This commit is contained in:
parent
28b331f4e8
commit
cc119297ba
|
@ -1,12 +1,11 @@
|
|||
import torch.nn
|
||||
import ldm.modules.diffusionmodules.openaimodel
|
||||
|
||||
import time
|
||||
from modules import script_callbacks, shared, devices
|
||||
|
||||
unet_options = []
|
||||
current_unet_option = None
|
||||
current_unet = None
|
||||
|
||||
original_forward = None # not used, only left temporarily for compatibility
|
||||
|
||||
def list_unets():
|
||||
new_unets = script_callbacks.list_unets_callback()
|
||||
|
@ -84,7 +83,10 @@ class SdUnet(torch.nn.Module):
|
|||
pass
|
||||
|
||||
|
||||
def create_unet_forward(original_forward):
|
||||
def UNetModel_forward(self, x, timesteps=None, context=None, *args, **kwargs):
|
||||
if current_unet is not None:
|
||||
return current_unet.forward(x, timesteps, context, *args, **kwargs)
|
||||
try:
|
||||
if current_unet is not None and shared.current_prompt != shared.skip_unet_prompt:
|
||||
if '[TRT]' in shared.opts.sd_unet and '<lora:' in shared.current_prompt:
|
||||
|
@ -98,4 +100,8 @@ def UNetModel_forward(self, x, timesteps=None, context=None, *args, **kwargs):
|
|||
shared.skip_unet_prompt = shared.current_prompt
|
||||
print('[UNet] Used', time.time() - start, 'seconds')
|
||||
|
||||
return ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui(self, x, timesteps, context, *args, **kwargs)
|
||||
|
||||
return original_forward(self, x, timesteps, context, *args, **kwargs)
|
||||
|
||||
return UNetModel_forward
|
||||
|
||||
|
|
Loading…
Reference in New Issue