add fused projections

This commit is contained in:
Vladimir Mandic
2024-01-02 09:41:15 -05:00
parent dc9778b701
commit f7cd47d96f
8 changed files with 60 additions and 10 deletions
+36
View File
@@ -130,6 +130,17 @@ def compile_torch(sd_model):
torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access
torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access
torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access
try:
torch._inductor.config.conv_1x1_as_mm = True # pylint: disable=protected-access
torch._inductor.config.coordinate_descent_tuning = True # pylint: disable=protected-access
torch._inductor.config.epilogue_fusion = False # pylint: disable=protected-access
torch._inductor.config.coordinate_descent_check_all_directions = True # pylint: disable=protected-access
torch._inductor.config.use_mixed_mm = True # pylint: disable=protected-access
# torch._inductor.config.force_fuse_int_mm_with_mul = True # pylint: disable=protected-access
except Exception as e:
shared.log.error(f"Torch inductor config error: {e}")
t0 = time.time()
if shared.opts.cuda_compile:
if shared.opts.cuda_compile and (not hasattr(sd_model, 'unet') or not hasattr(sd_model.unet, 'config')):
@@ -169,3 +180,28 @@ def compile_diffusers(sd_model):
else:
sd_model = compile_torch(sd_model)
return sd_model
def dynamic_quantization(sd_model):
try:
from torchao.quantization import quant_api
except Exception as e:
shared.log.error(f"Model dynamic quantization not supported: {e}")
return sd_model
def dynamic_quant_filter_fn(mod, *args): # pylint: disable=unused-argument
return (isinstance(mod, torch.nn.Linear) and mod.in_features > 16 and (mod.in_features, mod.out_features)
not in [(1280, 640), (1920, 1280), (1920, 640), (2048, 1280), (2048, 2560), (2560, 1280), (256, 128), (2816, 1280), (320, 640), (512, 1536), (512, 256), (512, 512), (640, 1280), (640, 1920), (640, 320), (640, 5120), (640, 640), (960, 320), (960, 640)])
def conv_filter_fn(mod, *args): # pylint: disable=unused-argument
return (isinstance(mod, torch.nn.Conv2d) and mod.kernel_size == (1, 1) and 128 in [mod.in_channels, mod.out_channels])
shared.log.info(f"Model dynamic quantization: pipeline={sd_model.__class__.__name__}")
try:
quant_api.swap_conv2d_1x1_to_linear(sd_model.unet, conv_filter_fn)
quant_api.swap_conv2d_1x1_to_linear(sd_model.vae, conv_filter_fn)
quant_api.apply_dynamic_quant(sd_model.unet, dynamic_quant_filter_fn)
quant_api.apply_dynamic_quant(sd_model.vae, dynamic_quant_filter_fn)
except Exception as e:
shared.log.error(f"Model dynamic quantization error: {e}")
return sd_model