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,101 +1,107 @@
|
|||
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
|
||||
|
||||
|
||||
def list_unets():
|
||||
new_unets = script_callbacks.list_unets_callback()
|
||||
|
||||
unet_options.clear()
|
||||
unet_options.extend(new_unets)
|
||||
|
||||
|
||||
def get_unet_option(option=None):
|
||||
option = option or shared.opts.sd_unet
|
||||
|
||||
if option == "None":
|
||||
return None
|
||||
|
||||
if option == "Automatic":
|
||||
name = shared.sd_model.sd_checkpoint_info.model_name
|
||||
|
||||
options = [x for x in unet_options if x.model_name == name]
|
||||
|
||||
option = options[0].label if options else "None"
|
||||
|
||||
return next(iter([x for x in unet_options if x.label == option]), None)
|
||||
|
||||
|
||||
def apply_unet(option=None):
|
||||
global current_unet_option
|
||||
global current_unet
|
||||
|
||||
new_option = get_unet_option(option)
|
||||
if new_option == current_unet_option:
|
||||
return
|
||||
|
||||
if current_unet is not None:
|
||||
print(f"Dectivating unet: {current_unet.option.label}")
|
||||
current_unet.deactivate()
|
||||
|
||||
current_unet_option = new_option
|
||||
if current_unet_option is None:
|
||||
current_unet = None
|
||||
|
||||
if not shared.sd_model.lowvram:
|
||||
shared.sd_model.model.diffusion_model.to(devices.device)
|
||||
|
||||
return
|
||||
|
||||
shared.sd_model.model.diffusion_model.to(devices.cpu)
|
||||
devices.torch_gc()
|
||||
|
||||
current_unet = current_unet_option.create_unet()
|
||||
current_unet.option = current_unet_option
|
||||
print(f"Activating unet: {current_unet.option.label}")
|
||||
current_unet.activate()
|
||||
|
||||
|
||||
class SdUnetOption:
|
||||
model_name = None
|
||||
"""name of related checkpoint - this option will be selected automatically for unet if the name of checkpoint matches this"""
|
||||
|
||||
label = None
|
||||
"""name of the unet in UI"""
|
||||
|
||||
def create_unet(self):
|
||||
"""returns SdUnet object to be used as a Unet instead of built-in unet when making pictures"""
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class SdUnet(torch.nn.Module):
|
||||
def forward(self, x, timesteps, context, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
|
||||
def activate(self):
|
||||
pass
|
||||
|
||||
def deactivate(self):
|
||||
pass
|
||||
|
||||
|
||||
def UNetModel_forward(self, x, timesteps=None, context=None, *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:
|
||||
raise Exception('LoRA unsupported in TRT UNet')
|
||||
f = current_unet.forward(x, timesteps, context, *args, **kwargs)
|
||||
return f
|
||||
except Exception as e:
|
||||
start = time.time()
|
||||
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 ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui(self, x, timesteps, context, *args, **kwargs)
|
||||
import torch.nn
|
||||
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()
|
||||
|
||||
unet_options.clear()
|
||||
unet_options.extend(new_unets)
|
||||
|
||||
|
||||
def get_unet_option(option=None):
|
||||
option = option or shared.opts.sd_unet
|
||||
|
||||
if option == "None":
|
||||
return None
|
||||
|
||||
if option == "Automatic":
|
||||
name = shared.sd_model.sd_checkpoint_info.model_name
|
||||
|
||||
options = [x for x in unet_options if x.model_name == name]
|
||||
|
||||
option = options[0].label if options else "None"
|
||||
|
||||
return next(iter([x for x in unet_options if x.label == option]), None)
|
||||
|
||||
|
||||
def apply_unet(option=None):
|
||||
global current_unet_option
|
||||
global current_unet
|
||||
|
||||
new_option = get_unet_option(option)
|
||||
if new_option == current_unet_option:
|
||||
return
|
||||
|
||||
if current_unet is not None:
|
||||
print(f"Dectivating unet: {current_unet.option.label}")
|
||||
current_unet.deactivate()
|
||||
|
||||
current_unet_option = new_option
|
||||
if current_unet_option is None:
|
||||
current_unet = None
|
||||
|
||||
if not shared.sd_model.lowvram:
|
||||
shared.sd_model.model.diffusion_model.to(devices.device)
|
||||
|
||||
return
|
||||
|
||||
shared.sd_model.model.diffusion_model.to(devices.cpu)
|
||||
devices.torch_gc()
|
||||
|
||||
current_unet = current_unet_option.create_unet()
|
||||
current_unet.option = current_unet_option
|
||||
print(f"Activating unet: {current_unet.option.label}")
|
||||
current_unet.activate()
|
||||
|
||||
|
||||
class SdUnetOption:
|
||||
model_name = None
|
||||
"""name of related checkpoint - this option will be selected automatically for unet if the name of checkpoint matches this"""
|
||||
|
||||
label = None
|
||||
"""name of the unet in UI"""
|
||||
|
||||
def create_unet(self):
|
||||
"""returns SdUnet object to be used as a Unet instead of built-in unet when making pictures"""
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class SdUnet(torch.nn.Module):
|
||||
def forward(self, x, timesteps, context, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
|
||||
def activate(self):
|
||||
pass
|
||||
|
||||
def deactivate(self):
|
||||
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:
|
||||
raise Exception('LoRA unsupported in TRT UNet')
|
||||
f = current_unet.forward(x, timesteps, context, *args, **kwargs)
|
||||
return f
|
||||
except Exception as e:
|
||||
start = time.time()
|
||||
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
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue