diffusers/tools

1#!/usr/bin/env python32import sys3 4filenames = sys.argv[1:]5 6MATCH_PATTERN_1 = "# Copied from transformers.models.bart.modeling_bart._make_causal_mask"7MATCH_PATTERN_2 = "def _make_causal_mask("8 9MATCH_PATTERN_1 = "# Copied from transformers.models.bart.modeling_bart.prepare_4d_attention_mask"10MATCH_PATTERN_2 = "def prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):"11 12END_MATCH_PATTERN_2 = ""13 14# MATCH_PATTERN_1 = "def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):"15#MATCH_PATTERN_2 = "# create causal mask"16 17# END_MATCH_PATTERN_2 = "def forward("18 19                                        20for filename in filenames:21    with open(filename, "r") as f:22        lines = f.readlines()23 24    new_lines = []25    is_in_del = False26    for i, line in enumerate(lines):27        if line.strip().lstrip() == MATCH_PATTERN_1 and i < len(lines) - 1 and lines[i + 1].strip().lstrip() == MATCH_PATTERN_2:28            print("suh")29            is_in_del = True30        elif line.strip().lstrip() == "" and i < len(lines) - 1 and lines[i + 1].strip().lstrip() == END_MATCH_PATTERN_2:31            is_in_del = False32 33        if not is_in_del:34            new_lines.append(line)35 36 37    with open(filename, "w") as f:38        f.writelines(new_lines)39