automatic/modules/hidiffusion/__init__.py

44 lines
2.7 KiB
Python

# Original: https://github.com/megvii-research/HiDiffusion
import time
from modules import shared
from modules.hidiffusion import hidiffusion
def apply(p, model_type):
if not shared.native:
return
if model_type not in ['sd', 'sdxl'] and p.hidiffusion:
shared.log.warning(f'HiDiffusion: class={shared.sd_model.__class__.__name__} not supported')
return
unapply()
if getattr(p, 'hidiffusion', False) is True:
t0 = time.time()
hidiffusion.is_aggressive_raunet = shared.opts.hidiffusion_steps > 0
hidiffusion.aggressive_step = shared.opts.hidiffusion_steps
if shared.opts.hidiffusion_t1 >= 0:
t1 = shared.opts.hidiffusion_t1
hidiffusion.switching_threshold_ratio_dict['sd15_1024']['T1_ratio'] = t1
hidiffusion.switching_threshold_ratio_dict['sd15_2048']['T1_ratio'] = t1
hidiffusion.switching_threshold_ratio_dict['sdxl_2048']['T1_ratio'] = t1
hidiffusion.switching_threshold_ratio_dict['sdxl_4096']['T1_ratio'] = t1
hidiffusion.switching_threshold_ratio_dict['sdxl_turbo_1024']['T1_ratio'] = t1
p.extra_generation_params['HiDiffusion Ratios'] = f'{shared.opts.hidiffusion_t1}/{shared.opts.hidiffusion_t2}'
if shared.opts.hidiffusion_t2 >= 0:
t2 =shared.opts.hidiffusion_t2
hidiffusion.switching_threshold_ratio_dict['sd15_1024']['T2_ratio'] = t2
hidiffusion.switching_threshold_ratio_dict['sd15_2048']['T2_ratio'] = t2
hidiffusion.switching_threshold_ratio_dict['sdxl_2048']['T2_ratio'] = t2
hidiffusion.switching_threshold_ratio_dict['sdxl_4096']['T2_ratio'] = t2
hidiffusion.switching_threshold_ratio_dict['sdxl_turbo_1024']['T2_ratio'] = t2
p.extra_generation_params['HiDiffusion Ratios'] = f'{shared.opts.hidiffusion_t1}/{shared.opts.hidiffusion_t2}'
hidiffusion.apply_hidiffusion(shared.sd_model, apply_raunet=shared.opts.hidiffusion_raunet, apply_window_attn=shared.opts.hidiffusion_attn, model_type=model_type)
p.extra_generation_params['HiDiffusion'] = f'{shared.opts.hidiffusion_raunet}/{shared.opts.hidiffusion_attn}/{shared.opts.hidiffusion_steps > 0}:{shared.opts.hidiffusion_steps}'
t1 = time.time()
shared.log.debug(f'HiDiffusion apply: raunet={shared.opts.hidiffusion_raunet} attn={shared.opts.hidiffusion_attn} aggressive={shared.opts.hidiffusion_steps > 0}:{shared.opts.hidiffusion_steps} t1={shared.opts.hidiffusion_t1} t2={shared.opts.hidiffusion_t2} time={t1-t0:.2f} type={shared.sd_model_type} width={p.width} height={p.height}')
def unapply():
if hasattr(shared.sd_model, "unet"):
hidiffusion.remove_hidiffusion(shared.sd_model)