cofob
/
Eris_api_tensor_patch
forked from Hitmare/Eris_api_tensor_patch
1
0
Fork 0
Code Pull Requests Activity

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:
Hitmare 2023-12-23 15:03:35 +00:00
parent 28b331f4e8
commit cc119297ba
1 changed files with 107 additions and 101 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
TRT-Patch/modules/sd_unet.py
View File

@ -1,12 +1,11 @@
import torch.nn import torch.nn
import ldm.modules.diffusionmodules.openaimodel
import time import time
from modules import script_callbacks, shared, devices from modules import script_callbacks, shared, devices
unet_options = [] unet_options = []
current_unet_option = None current_unet_option = None
current_unet = None current_unet = None
original_forward = None # not used, only left temporarily for compatibility
def list_unets(): def list_unets():
new_unets = script_callbacks.list_unets_callback() new_unets = script_callbacks.list_unets_callback()
@ -84,18 +83,25 @@ class SdUnet(torch.nn.Module):
pass pass
def UNetModel_forward(self, x, timesteps=None, context=None, *args, **kwargs): def create_unet_forward(original_forward):
try: def UNetModel_forward(self, x, timesteps=None, context=None, *args, **kwargs):
if current_unet is not None and shared.current_prompt != shared.skip_unet_prompt: if current_unet is not None:
if '[TRT]' in shared.opts.sd_unet and '<lora:' in shared.current_prompt: return current_unet.forward(x, timesteps, context, *args, **kwargs)
raise Exception('LoRA unsupported in TRT UNet') try:
f = current_unet.forward(x, timesteps, context, *args, **kwargs) if current_unet is not None and shared.current_prompt != shared.skip_unet_prompt:
return f if '[TRT]' in shared.opts.sd_unet and '<lora:' in shared.current_prompt:
except Exception as e: raise Exception('LoRA unsupported in TRT UNet')
start = time.time() f = current_unet.forward(x, timesteps, context, *args, **kwargs)
print('[UNet] Skipping TRT UNet for this request:', e, '-', shared.current_prompt) return f
shared.sd_model.model.diffusion_model.to(devices.device) except Exception as e:
shared.skip_unet_prompt = shared.current_prompt start = time.time()
print('[UNet] Used', time.time() - start, 'seconds') print('[UNet] Skipping TRT UNet for this request:', e, '-', shared.current_prompt)
shared.sd_model.model.diffusion_model.to(devices.device)
shared.skip_unet_prompt = shared.current_prompt
print('[UNet] Used', time.time() - start, 'seconds')
return original_forward(self, x, timesteps, context, *args, **kwargs)
return UNetModel_forward
return ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui(self, x, timesteps, context, *args, **kwargs)