diff --git a/tools/pot/openvino/tools/pot/graph/special_patterns.py b/tools/pot/openvino/tools/pot/graph/special_patterns.py index 2ac1eac6962..d0810a99622 100644 --- a/tools/pot/openvino/tools/pot/graph/special_patterns.py +++ b/tools/pot/openvino/tools/pot/graph/special_patterns.py @@ -293,6 +293,21 @@ def create_softmax_reshape_transpose_matmul_pattern(): return pattern.set_name('softmax_reshape_transpose_matmul').pattern +# Stable diffusion UNet +@registry_ignore_patterns('blocks') +def create_stable_diffusion_pattern(): + pattern = PatternBuilder() + pattern_2 = PatternBuilder() + softmax_out = pattern.append_single_op('SoftMax', 'softmax').get_last_node() + pattern_2.append_single_op('Reshape', 'reshape1') + pattern_2.append_single_op('Transpose', 'transpose') + transp_out = pattern_2.append_single_op('Reshape', 'reshape2').get_last_node() + pattern.pattern['nodes'] += pattern_2.pattern['nodes'] + pattern.pattern['edges'] += pattern_2.pattern['edges'] + pattern.insert_single_op([transp_out, softmax_out], None, 'MatMul', 'matmul') + return pattern.set_name('stable_diffusion').pattern + + @registry_ignore_patterns('blocks') def create_softmax_reshape_transpose_gather_matmul_pattern(): pattern = PatternBuilder()