├── .gitignore ├── docs ├── add.png ├── delete.png ├── workflow.png └── replacement.png ├── prompt2prompt ├── Roboto-Regular.ttf ├── __pycache__ │ ├── ptp_utils.cpython-39.pyc │ └── attn_control.cpython-39.pyc ├── attn_control.py └── ptp_utils.py ├── utils ├── __pycache__ │ ├── utils.cpython-39.pyc │ └── converter.cpython-39.pyc ├── utils.py └── converter.py ├── suppresseot ├── __pycache__ │ ├── attn_loss.cpython-39.pyc │ ├── wo_utils.cpython-39.pyc │ └── run_and_display.cpython-39.pyc ├── wo_utils.py ├── attn_loss.py └── run_and_display.py ├── auffusion ├── __pycache__ │ └── auffusion_pipeline.cpython-39.pyc └── auffusion_pipeline.py ├── audio_examples └── input_audios │ ├── A woman is giving a speech.wav │ └── After a gunshot, there was a burst of dog barking.wav ├── null_text_inversion ├── __pycache__ │ └── null_text_inversion.cpython-39.pyc └── null_text_inversion.py ├── eval ├── clap_score │ └── cal_clap_score.py └── configs │ └── audiocaps │ ├── edit_replace.csv │ ├── edit_delete.csv │ └── edit_add.csv ├── main.sh ├── requirements.txt ├── main.py ├── readme.md └── LICENSE /.gitignore: -------------------------------------------------------------------------------- 1 | audio_examples/output_audios/* 2 | ckpt/* -------------------------------------------------------------------------------- /docs/add.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/docs/add.png -------------------------------------------------------------------------------- /docs/delete.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/docs/delete.png -------------------------------------------------------------------------------- /docs/workflow.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/docs/workflow.png -------------------------------------------------------------------------------- /docs/replacement.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/docs/replacement.png -------------------------------------------------------------------------------- /prompt2prompt/Roboto-Regular.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/prompt2prompt/Roboto-Regular.ttf -------------------------------------------------------------------------------- /utils/__pycache__/utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/utils/__pycache__/utils.cpython-39.pyc -------------------------------------------------------------------------------- /utils/__pycache__/converter.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/utils/__pycache__/converter.cpython-39.pyc -------------------------------------------------------------------------------- /suppresseot/__pycache__/attn_loss.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/suppresseot/__pycache__/attn_loss.cpython-39.pyc -------------------------------------------------------------------------------- /suppresseot/__pycache__/wo_utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/suppresseot/__pycache__/wo_utils.cpython-39.pyc -------------------------------------------------------------------------------- /prompt2prompt/__pycache__/ptp_utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/prompt2prompt/__pycache__/ptp_utils.cpython-39.pyc -------------------------------------------------------------------------------- /prompt2prompt/__pycache__/attn_control.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/prompt2prompt/__pycache__/attn_control.cpython-39.pyc -------------------------------------------------------------------------------- /suppresseot/__pycache__/run_and_display.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/suppresseot/__pycache__/run_and_display.cpython-39.pyc -------------------------------------------------------------------------------- /auffusion/__pycache__/auffusion_pipeline.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/auffusion/__pycache__/auffusion_pipeline.cpython-39.pyc -------------------------------------------------------------------------------- /audio_examples/input_audios/A woman is giving a speech.wav: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/audio_examples/input_audios/A woman is giving a speech.wav -------------------------------------------------------------------------------- /null_text_inversion/__pycache__/null_text_inversion.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/null_text_inversion/__pycache__/null_text_inversion.cpython-39.pyc -------------------------------------------------------------------------------- /audio_examples/input_audios/After a gunshot, there was a burst of dog barking.wav: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NKU-HLT/AudioEditor/HEAD/audio_examples/input_audios/After a gunshot, there was a burst of dog barking.wav -------------------------------------------------------------------------------- /eval/clap_score/cal_clap_score.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from msclap import CLAP 3 | 4 | # Load model (Choose between versions '2022' or '2023') 5 | # The model weight will be downloaded automatically if `model_fp` is not specified 6 | clap_model = CLAP(version = '2023', model_fp = 'CLAP_weights_2023.pth', use_cuda = False) 7 | 8 | # Extract text embeddings 9 | text_embeddings = clap_model.get_text_embeddings(["A baby and a woman laugh"]) 10 | 11 | # Extract audio embeddings 12 | audio_embeddings = clap_model.get_audio_embeddings(["A baby and a woman laugh.wav"]) 13 | 14 | # cosine_similarity = torch.nn.functional.cosine_similarity(audio_embeddings, text_embeddings) 15 | # print("Cosine Similarity:", cosine_similarity.item()) -------------------------------------------------------------------------------- /main.sh: -------------------------------------------------------------------------------- 1 | export HF_ENDPOINT="https://hf-mirror.com" 2 | export HF_HUB_ENABLE_HF_TRANSFER=1 3 | 4 | CUDA_VISIBLE_DEVICES='0' 5 | 6 | ### delete 7 | python main.py --prompt "After a gunshot, there was a burst of dog barking" \ 8 | --audio_path "audio_examples/input_audios/After a gunshot, there was a burst of dog barking.wav" \ 9 | --token_indices "[[10,11]]" \ 10 | --alpha "[1.,]" --cross_retain_steps "[.2,]" 11 | 12 | ### replace 13 | python main.py --prompt "After a thunder, there was a burst of dog barking" \ 14 | --audio_path "audio_examples/input_audios/After a gunshot, there was a burst of dog barking.wav" \ 15 | --token_indices "[[3]]" \ 16 | --alpha "[-0.001,]" --cross_retain_steps "[.2,]" 17 | 18 | ### add 19 | python main.py --prompt "A woman is giving a speech amid applause" \ 20 | --audio_path "audio_examples/input_audios/A woman is giving a speech.wav" \ 21 | --token_indices "[[7,8]]" \ 22 | --alpha "[-0.001,]" --cross_retain_steps "[.2,]" -------------------------------------------------------------------------------- /suppresseot/wo_utils.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import random 4 | from scipy.spatial.distance import cdist 5 | 6 | CALC_SIMILARITY = False 7 | 8 | def punish_wight(wo_batch, latent_size, alpha, method): 9 | if method == 'weight': 10 | wo_batch *= alpha 11 | elif method in ['alpha', 'beta', 'delete', 'soft-weight']: 12 | u, s, vh = torch.linalg.svd(wo_batch) 13 | u = u[:,:latent_size] 14 | zero_idx = int(latent_size * alpha) 15 | if method == 'alpha': 16 | s[:zero_idx] = 0 17 | elif method == 'beta': 18 | s[zero_idx:] = 0 19 | elif method == 'delete': 20 | s = s[zero_idx:] if zero_idx < latent_size else torch.zeros(latent_size).to(s.device) 21 | u = u[:, zero_idx:] if zero_idx < latent_size else u 22 | vh = vh[zero_idx:, :] if zero_idx < latent_size else vh 23 | elif method == 'soft-weight': 24 | if CALC_SIMILARITY: 25 | _s = s.clone() 26 | _s[zero_idx:] = 0 27 | _wo_batch = u @ torch.diag(_s) @ vh 28 | dist = cdist(wo_batch[:,0].unsqueeze(0).cpu(), _wo_batch[:,0].unsqueeze(0).cpu(), metric='cosine') 29 | print(f'The distance between the word embedding before and after the punishment: {dist}') 30 | if alpha == -.001: 31 | s *= (torch.exp(-.001 * s) * 1.2) # add $ replace 32 | else: 33 | s *= torch.exp(-alpha*s) # delete 34 | 35 | wo_batch = u @ torch.diag(s) @ vh 36 | else: 37 | raise ValueError('Unsupported method') 38 | return wo_batch 39 | 40 | def woword_eot_context(context, token_indices, alpha, method, n): 41 | for i, batch in enumerate(context): 42 | indices = token_indices + [num for num in range(n-1, 77)] 43 | wo_batch = batch[indices] 44 | wo_batch = punish_wight(wo_batch.T, len(indices), alpha, method).T.float() # .float() 45 | batch[indices] = wo_batch 46 | return context 47 | 48 | def woword_reweight(attn, token_indices, alpha): 49 | # if attn.shape[1] > 32 ** 2: # avoid memory overhead 50 | # return attn 51 | latent_size = int(attn.shape[1]**0.5) 52 | assert latent_size**2 == attn.shape[1] 53 | for head_attn in attn: 54 | for indice in token_indices: 55 | wo_attn = head_attn[:, indice].reshape(latent_size, latent_size) 56 | wo_attn *= alpha # same as Reweight of P2P 57 | head_attn[:, indice] = wo_attn.reshape(latent_size**2) 58 | return attn -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | accelerate==0.32.0 2 | aiohttp==3.9.5 3 | aiosignal==1.3.1 4 | asttokens==2.4.1 5 | async-timeout==4.0.3 6 | attrs==23.2.0 7 | audioread==3.0.1 8 | certifi==2024.7.4 9 | cffi==1.16.0 10 | charset-normalizer==3.3.2 11 | cmake==3.30.0 12 | comm==0.2.2 13 | contourpy==1.2.1 14 | cycler==0.12.1 15 | datasets==2.20.0 16 | debugpy==1.8.2 17 | decorator==5.1.1 18 | diffusers==0.18.2 19 | dill==0.3.8 20 | einops==0.8.0 21 | exceptiongroup==1.2.1 22 | executing==2.0.1 23 | filelock==3.15.4 24 | fonttools==4.53.0 25 | frozenlist==1.4.1 26 | fsspec==2024.5.0 27 | hf_transfer==0.1.6 28 | huggingface-hub==0.23.4 29 | idna==3.7 30 | importlib_metadata==8.0.0 31 | importlib_resources==6.4.0 32 | ipykernel==6.29.5 33 | ipython==8.18.1 34 | jedi==0.19.1 35 | Jinja2==3.1.4 36 | joblib==1.4.2 37 | jupyter_client==8.6.2 38 | jupyter_core==5.7.2 39 | kiwisolver==1.4.5 40 | librosa==0.9.2 41 | lit==18.1.8 42 | llvmlite==0.43.0 43 | MarkupSafe==2.1.5 44 | matplotlib==3.9.1 45 | matplotlib-inline==0.1.7 46 | mpmath==1.3.0 47 | multidict==6.0.5 48 | multiprocess==0.70.16 49 | mypy-extensions==1.0.0 50 | nest-asyncio==1.6.0 51 | networkx==3.2.1 52 | numba==0.60.0 53 | numpy==1.26.4 54 | nvidia-cublas-cu11==11.10.3.66 55 | nvidia-cuda-cupti-cu11==11.7.101 56 | nvidia-cuda-nvrtc-cu11==11.7.99 57 | nvidia-cuda-runtime-cu11==11.7.99 58 | nvidia-cudnn-cu11==8.5.0.96 59 | nvidia-cufft-cu11==10.9.0.58 60 | nvidia-curand-cu11==10.2.10.91 61 | nvidia-cusolver-cu11==11.4.0.1 62 | nvidia-cusparse-cu11==11.7.4.91 63 | nvidia-nccl-cu11==2.14.3 64 | nvidia-nvtx-cu11==11.7.91 65 | opencv-python==4.10.0.84 66 | packaging==24.1 67 | pandas==2.2.2 68 | parso==0.8.4 69 | pexpect==4.9.0 70 | pillow==10.4.0 71 | platformdirs==4.2.2 72 | pooch==1.8.2 73 | prompt_toolkit==3.0.47 74 | psutil==6.0.0 75 | ptyprocess==0.7.0 76 | pure-eval==0.2.2 77 | pyarrow==16.1.0 78 | pyarrow-hotfix==0.6 79 | pycparser==2.22 80 | Pygments==2.18.0 81 | pyparsing==3.1.2 82 | pyre-extensions==0.0.29 83 | python-dateutil==2.9.0.post0 84 | pytz==2024.1 85 | PyYAML==6.0.1 86 | pyzmq==26.0.3 87 | regex==2024.5.15 88 | requests==2.32.3 89 | resampy==0.4.3 90 | safetensors==0.4.3 91 | scikit-learn==1.5.1 92 | scipy==1.13.1 93 | six==1.16.0 94 | soundfile==0.12.1 95 | stack-data==0.6.3 96 | sympy==1.12.1 97 | threadpoolctl==3.5.0 98 | tokenizers==0.13.3 99 | torch==2.0.1 100 | torchaudio==2.0.2 101 | torchvision==0.15.2 102 | tornado==6.4.1 103 | tqdm==4.66.4 104 | traitlets==5.14.3 105 | transformers==4.33.3 106 | triton==2.0.0 107 | typing-inspect==0.9.0 108 | typing_extensions==4.12.2 109 | tzdata==2024.1 110 | urllib3==2.2.2 111 | wcwidth==0.2.13 112 | xformers==0.0.20 113 | xxhash==3.4.1 114 | yarl==1.9.4 115 | zipp==3.19.2 116 | -------------------------------------------------------------------------------- /suppresseot/attn_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.nn import functional as F 4 | 5 | class Loss(torch.nn.Module): 6 | def __init__(self, loss_type='mse'): 7 | super(Loss, self).__init__() 8 | 9 | self.loss_type = loss_type 10 | 11 | self.loss_func = { 12 | 'mse': torch.nn.MSELoss, 13 | 'cosine': torch.nn.CosineSimilarity, 14 | 'mae': torch.nn.L1Loss 15 | }[loss_type]() 16 | 17 | def forward(self, x, y): 18 | if self.loss_type == "cosine": 19 | return 1. - self.loss_func(x, y) 20 | 21 | return self.loss_func(x, y) 22 | 23 | class AttnLoss(nn.Module): 24 | def __init__(self, device, attn_loss_type, n, token_indices, lambda_retain=1., lambda_erase=-1., lambda_self_retain=1., lambda_self_erase=-1.): 25 | super(AttnLoss, self).__init__() 26 | self.device = device 27 | self.prompt_n = n 28 | self.token_indices = token_indices 29 | 30 | self.lambda_retain = lambda_retain 31 | self.lambda_erase = lambda_erase 32 | self.lambda_self_retain = lambda_self_retain 33 | self.lambda_self_erase = lambda_self_erase 34 | 35 | self.retain_loss = Loss(attn_loss_type) 36 | self.erase_loss = Loss(attn_loss_type) 37 | self.self_retain_loss = Loss(attn_loss_type) 38 | self.self_erase_loss = Loss(attn_loss_type) 39 | 40 | def calc_mask(self, attn, threshold=.85): 41 | mask = [] 42 | for i in [num for num in range(1, self.prompt_n-1)]: 43 | _attn = attn[:,:,i].clone() 44 | _attn = 255 * _attn / _attn.max() 45 | _attn = F.interpolate(_attn.unsqueeze(0).unsqueeze(0), size=(256, 256), mode='bilinear') 46 | if i in self.token_indices: 47 | _threshold = threshold 48 | else: 49 | _threshold = threshold + .1 50 | _attn[_attn >= _attn.max() * _threshold] = 255 51 | _attn[_attn < _attn.max() * _threshold] = 0 52 | _attn = F.interpolate(_attn, size=attn.shape[:2], mode='bilinear') 53 | mask += [_attn.squeeze(0).squeeze(0)] 54 | return mask 55 | 56 | def calc_retain_loss(self, attn, attn_erase): 57 | loss = .0 58 | for i in [num for num in range(1, self.prompt_n-1)]: 59 | if i in self.token_indices: 60 | continue 61 | loss += self.retain_loss(attn[:,:,i], attn_erase[:,:,i]) 62 | return loss 63 | 64 | def calc_erase_loss(self, attn, attn_erase): 65 | loss = .0 66 | for i in self.token_indices: 67 | loss += self.erase_loss(attn[:,:,i], attn_erase[:,:,i]) 68 | return loss 69 | 70 | def calc_self_retain_loss(self, self_attn, self_attn_erase, mask): 71 | loss = .0 72 | h, w = mask[0].shape 73 | for i in [num for num in range(1, self.prompt_n-1)]: 74 | if i in self.token_indices: 75 | continue 76 | for j, m in enumerate(mask[i-1].reshape(h*w)): 77 | if m > 0: 78 | loss += self.self_retain_loss(self_attn[:,:,j].view(-1).unsqueeze(0), 79 | self_attn_erase[:,:,j].view(-1).unsqueeze(0)) 80 | return loss 81 | 82 | def calc_self_erase_loss(self, self_attn, self_attn_erase, mask): 83 | loss = .0 84 | h, w = mask[0].shape 85 | for i in self.token_indices: 86 | for j, m in enumerate(mask[i-1].reshape(h*w)): 87 | if m > 0: 88 | loss += self.self_erase_loss(self_attn[:,:,j].view(-1).unsqueeze(0), 89 | self_attn_erase[:,:,j].view(-1).unsqueeze(0)) 90 | return loss 91 | 92 | def forward(self, attn, attn_erase, self_attn, self_attn_erase): 93 | attn, attn_erase, self_attn, self_attn_erase \ 94 | = attn.to(torch.double), attn_erase.to(torch.double), self_attn.to(torch.double), self_attn_erase.to(torch.double) 95 | attn_loss = .0 96 | 97 | if self.lambda_self_retain or self.lambda_self_erase: 98 | mask = self.calc_mask(attn) 99 | 100 | h, w, seq_len = attn.shape 101 | attn = attn.reshape(h*w, seq_len).unsqueeze(0) 102 | attn_erase = attn_erase.reshape(h*w, seq_len).unsqueeze(0) 103 | 104 | if self.lambda_retain: 105 | attn_loss += self.lambda_retain * self.calc_retain_loss(attn, attn_erase) 106 | 107 | if self.lambda_erase: 108 | attn_loss += self.lambda_erase * self.calc_erase_loss(attn, attn_erase) 109 | 110 | if self.lambda_self_retain: 111 | attn_loss += self.lambda_self_retain * self.calc_self_retain_loss(self_attn, self_attn_erase, mask) 112 | 113 | if self.lambda_self_erase: 114 | attn_loss += self.lambda_self_erase * self.calc_self_erase_loss(self_attn, self_attn_erase, mask) 115 | 116 | loss = attn_loss 117 | return loss -------------------------------------------------------------------------------- /eval/configs/audiocaps/edit_replace.csv: -------------------------------------------------------------------------------- 1 | Uid,Caption,edit_Caption 2 | 5er9MoTAogc_10000_20000,A baby crying loudly,A baby laughing loudly 3 | z3gvS28om7o_180000_190000,A woman giving a speech,A man giving a speech 4 | qjwDugPZMjA_210000_220000,A crowd of people applauding,A crowd of people laughing 5 | P4Rf3_mKc4s_30000_40000,A woman is performing a speech,A man is performing a speech 6 | USle1qAiOdA_0_10000,A cat meowing twice,A dog barking twice 7 | DBjQUbJNHwc_30000_40000,A woman is giving a speech,A man is giving a speech 8 | xzQ9Q4UKOjY_90000_100000,A baby is crying,A baby is laughing 9 | zvLcOfhGcGo_0_10000,A person snoring,A person applause 10 | SaUDkGYRZHE_0_10000,A man laughing,A man crying 11 | GDMI4uJfDVw_10000_20000,A woman speaking,A woman crying 12 | UZqMaHTN6SE_0_10000,A telephone rings several times,A sirens rings several times 13 | zpMwBUIuxgY_360000_370000,Typing on a computer keyboard,Typing on a computer Piano 14 | z4omk5iedYc_80000_90000,A toilet is flushed,A engine is flushed 15 | ypY1zYRcfIc_100000_110000,A toilet flushing,A thunder flushing 16 | jEnNUSbVz3c_0_10000,Whistling a tune,Piano a tune 17 | 5p9XLBtfesc_330000_340000,A man speaks followed by applause,A man speaks followed by sirens 18 | Ekn0Q03AFN0_30000_40000,A woman speeches,A woman laughing 19 | zxvv9NC86uI_30000_40000,A young woman talking,A young woman laughing 20 | pvy7D5MvdIE_20000_30000,A man speaks followed by a crowd applauding,A man speaks followed by a crowd sirens 21 | 4kJQOM9BN4g_28000_38000,An engine is accelerating,An thunder is accelerating 22 | Mcvgo80QdlI_10000_20000,A person is snoring steadily,A person is crying steadily 23 | zssgRs7HUmM_30000_40000,A clock ticking,A engine ticking 24 | zD2KfYF4e2Q_10000_20000,Musical whistling,Piano whistling 25 | 8aNj2ItTaIs_180000_190000,Wind blows hard,thunder blows hard 26 | 78ER5FzxYdQ_30000_40000,Birds singing outdoors,dog barking outdoors 27 | yZpTj9YT-IQ_30000_40000,A man speaking into a microphone,A man laughing into a microphone 28 | yc6NJXKEIEk_30000_40000,An engine revving,An sirens revving 29 | Jzc66pasXSg_110000_120000,Repetitive snoring,Repetitive laughing 30 | 4wjKtAs9aMQ_29000_39000,An engine is idling,An thunder is idling 31 | ltQBJIueFT0_30000_40000,A male voice speaking,A male voice laughing 32 | 4Jo3aJriewE_30000_40000,A woman delivers a speech,A man delivers a speech 33 | NFnPqqOeG_A_280000_290000,Large bells ringing,Large thunder ringing 34 | z-3jiJkrDQE_10000_20000,A person is typing on a computer keyboard,A person is typing on a computer Piano 35 | xzQ9Q4UKOjY_90000_100000,A baby is crying,A baby is laughing 36 | qyfZgjNoJlE_0_10000,Continuous laughter,Continuous crying 37 | yvZWpypX7_Y_30000_40000,A woman speaking continuously,A woman applause continuously 38 | eK-tTmyMlgI_0_10000,Birds are chirping,dogs are chirping 39 | yizOr1hPxc8_70000_80000,A helicopter engine running idle,A car engine running idle 40 | 1MUcHcqDJfo_23000_33000,Car engine idling,Helicopter engine idling 41 | C9DKBh_c4xA_30000_40000,An elderly man speaking,An young man speaking 42 | ApqTL3dFTdM_310000_320000,A high pitched engine idling,A high pitched dog idling 43 | AhHW1jNf_cw_30000_40000,A man speaking followed by rustling,A man speaking followed by thunder 44 | twOUIlaeVm8_30000_40000,A baby crying continuously,A baby laughing continuously 45 | K6udqujyKFo_20000_30000,A person laughs,A person applause 46 | ikSxpwLxwKk_200000_210000,Adult female speaking,Little female speaking 47 | 9wFPkAuiOuI_280000_290000,Two men speaking,Two men laughing 48 | ztI9tn8Lfnk_9000_19000,Clicking on a computer keyboard,Clicking on a computer guitar 49 | 91EvnH5ln4Q_30000_40000,A female speaking,A female typing 50 | v-RSiQKNgM8_30000_40000,A sewing machine operating,A sewing engine operating 51 | sLPAcgyETPM_30000_40000,Female speaking,Female laughing 52 | pmeh7ocxLP4_0_10000,Continuous snoring,Continuous raining 53 | DBjQUbJNHwc_30000_40000,A woman is giving a speech,A woman is giving a piano 54 | 95vFDbFrq5g_30000_40000,A male speaking,A male typing 55 | zNLca6Dh6-k_120000_130000,Vibrations from a sewing machine,Vibrations from a sewing engine 56 | zO1h2bKYe58_140000_150000,An idle vehicle engine running,An idle helicopter engine running 57 | jcf9Smq-0Ps_20000_30000,Toilet flushing,Thunder flushing 58 | zNW0z-SHXTM_70000_80000,Humming of an idling engine,Humming of an idling thunder 59 | o7Nra5Cw_jQ_20000_30000,A telephone bell rings,A heavy sirens rings 60 | rGT-FW1kadU_30000_40000,"Heavy, continuous wind","Heavy, continuous whistles" 61 | ui4PHURhahA_100000_110000,A duck quacking,A dog barking 62 | BDDxfr4xf0k_60000_70000,A helicopter running,A car running 63 | 3y-3xiz2U5M_110000_120000,A power tool is in use,A piano kit is in use 64 | -ftxbusleO8_30000_40000,An engine running,An thunder running 65 | EGve_FMNOwE_30000_40000,A woman is speaking,A woman is laughing 66 | IEI10aLmHXM_30000_40000,A man is giving a speech,A man is giving a applause 67 | yWOeusquiJk_20000_30000,A person is snoring,A person is speaking 68 | U3COtQZORMU_20000_30000,A loud engine idling,A loud thunder idling 69 | ozQ428JESb8_110000_120000,A young girl giving a speech,A young girl giving a piano 70 | zNW0z-SHXTM_70000_80000,Humming of an idling engine,Humming of an dog barking 71 | z_zkX-8Kf3Q_0_10000,A crowd applauds,A crowd sirens 72 | gcF9BFScQHU_170000_180000,Multiple bells ringing,Multiple thunder ringing 73 | q4yBn1ECb5U_50000_60000,A toilet flushes,A sirens flushes 74 | yjLvzh0s1Ug_0_10000,A vehicle engine revving several times,A helicopter engine revving several times 75 | zNW0z-SHXTM_70000_80000,Humming of an idling engine,Humming of an idling cat 76 | z2QsWDZMXd0_30000_40000,Ticking of a clock,Ticking of a bell 77 | zRiGmSlWjok_200000_210000,"Loud, continuous applause","Loud, continuous thunder" 78 | ydO6ZIcghzU_0_10000,A person whistling,A person shouting 79 | ApMyWG-5x6M_20000_30000,An engine idles loudly,An car idles loudly 80 | d03ZdivCSKM_30000_40000,A baby cries continuously,A baby laughing continuously 81 | 8As1owPi19k_490000_500000,A man speaking then whistling,A man speaking then raining 82 | zf9Uk3HjD4k_20000_30000,A telephone ringing,A motorbike ringing 83 | Qtdw1gRUsLM_560000_570000,Several large bells ring,Several large clock ring 84 | xFndh2QPZG0_30000_40000,An engine running consistently,An dog running consistently 85 | zssgRs7HUmM_30000_40000,A clock ticking,A bell ticking 86 | UuST6ut8oEY_210000_220000,"A long, loud burp","A long, loud explosion" 87 | h5DUE-hRb5Y_30000_40000,Someone is whistling,Someone is laughing 88 | yPCgybma0Jc_310000_320000,Food frying in a pan,Food frying in a mouth 89 | v-RSiQKNgM8_30000_40000,A sewing machine operating,A sewing gunshot operating 90 | GEBMPrzW0NI_30000_40000,A man is speaking,A man is crying 91 | ns6btoVV_pA_20000_30000,A motorbike engine running idle,A helicopter engine running idle 92 | -------------------------------------------------------------------------------- /prompt2prompt/attn_control.py: -------------------------------------------------------------------------------- 1 | import abc 2 | 3 | import numpy as np 4 | import torch 5 | from IPython.display import display 6 | from PIL import Image 7 | from typing import Union, Tuple, List, Dict, Optional 8 | import matplotlib.pyplot as plt 9 | 10 | from prompt2prompt import ptp_utils 11 | 12 | LOW_RESOURCE = True 13 | NUM_DDIM_STEPS = 100 14 | 15 | class EmptyControl: 16 | 17 | def step_callback(self, x_t): 18 | return x_t 19 | 20 | def between_steps(self): 21 | return 22 | 23 | def __call__(self, attn, is_cross: bool, place_in_unet: str): 24 | return attn 25 | 26 | class AttentionControl(abc.ABC): 27 | 28 | def step_callback(self, x_t): 29 | return x_t 30 | 31 | def between_steps(self): 32 | return 33 | 34 | @property 35 | def num_uncond_att_layers(self): 36 | return self.num_att_layers if LOW_RESOURCE else 0 37 | 38 | @abc.abstractmethod 39 | def forward(self, attn, is_cross: bool, place_in_unet: str): 40 | raise NotImplementedError 41 | 42 | def __call__(self, attn, is_cross: bool, place_in_unet: str): 43 | if self.cur_att_layer >= self.num_uncond_att_layers: 44 | if LOW_RESOURCE: 45 | attn = self.forward(attn, is_cross, place_in_unet) 46 | else: 47 | h = attn.shape[0] 48 | attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet) 49 | self.cur_att_layer += 1 50 | if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers: 51 | self.cur_att_layer = 0 52 | self.cur_step += 1 53 | self.between_steps() 54 | return attn 55 | 56 | def reset(self): 57 | self.cur_step = 0 58 | self.cur_att_layer = 0 59 | 60 | def __init__(self): 61 | self.cur_step = 0 62 | self.num_att_layers = -1 63 | self.cur_att_layer = 0 64 | 65 | 66 | class SpatialReplace(EmptyControl): 67 | 68 | def step_callback(self, x_t): 69 | if self.cur_step < self.stop_inject: 70 | b = x_t.shape[0] 71 | x_t = x_t[:1].expand(b, *x_t.shape[1:]) 72 | return x_t 73 | 74 | def __init__(self, stop_inject: float): 75 | super(SpatialReplace, self).__init__() 76 | self.stop_inject = int((1 - stop_inject) * NUM_DDIM_STEPS) 77 | 78 | 79 | class AttentionStore(AttentionControl): 80 | 81 | @staticmethod 82 | def get_empty_store(): 83 | return {"down_cross": [], "mid_cross": [], "up_cross": [], 84 | "down_self": [], "mid_self": [], "up_self": []} 85 | 86 | def forward(self, attn, is_cross: bool, place_in_unet: str): 87 | key = f"{place_in_unet}_{'cross' if is_cross else 'self'}" 88 | if attn.shape[1] <= 32 ** 2: # avoid memory overhead 89 | self.step_store[key].append(attn) 90 | return attn 91 | 92 | def between_steps(self): 93 | if len(self.attention_store) == 0: 94 | self.attention_store = self.step_store 95 | else: 96 | for key in self.attention_store: 97 | for i in range(len(self.attention_store[key])): 98 | self.attention_store[key][i] += self.step_store[key][i] 99 | self.step_store = self.get_empty_store() 100 | 101 | def get_average_attention(self): 102 | average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in 103 | self.attention_store} 104 | return average_attention 105 | 106 | def reset(self): 107 | super(AttentionStore, self).reset() 108 | self.step_store = self.get_empty_store() 109 | self.attention_store = {} 110 | 111 | def __init__(self, token_indices: List[int], alpha: float, method: str, cross_retain_steps: float, n: int, iter_each_step: int, max_step_to_erase: int, 112 | lambda_retain=1, lambda_erase=-.5, lambda_self_retain=1, lambda_self_erase=-.5): 113 | super(AttentionStore, self).__init__() 114 | self.step_store = self.get_empty_store() 115 | self.attention_store = {} 116 | self.baseline = True 117 | # for suppression content 118 | self.ddim_inv = False 119 | self.token_indices = token_indices 120 | self.uncond = True 121 | self.alpha = alpha 122 | self.method = method # default: 'soft-weight' 123 | self.i = None 124 | self.cross_retain_steps = cross_retain_steps * NUM_DDIM_STEPS 125 | self.n = n 126 | self.text_embeddings_erase = None 127 | self.iter_each_step = iter_each_step 128 | self.MAX_STEP_TO_ERASE = max_step_to_erase 129 | # lambds of loss 130 | self.lambda_retain = lambda_retain 131 | self.lambda_erase = lambda_erase 132 | self.lambda_self_retain = lambda_self_retain 133 | self.lambda_self_erase = lambda_self_erase 134 | 135 | 136 | def aggregate_attention(prompts, attention_store: AttentionStore, res: List[int], from_where: List[str], is_cross: bool, select: int): 137 | out = [] 138 | attention_maps = attention_store.get_average_attention() 139 | num_pixels = res[0] * res[1] 140 | for location in from_where: 141 | for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]: 142 | if item.shape[1] == num_pixels: 143 | cross_maps = item.reshape(len(prompts), -1, res[0], res[1], item.shape[-1])[select] 144 | out.append(cross_maps) 145 | out = torch.cat(out, dim=0) 146 | out = out.sum(0) / out.shape[0] 147 | return out.cpu() 148 | 149 | def show_cross_attention(stable, prompts, attention_store: AttentionStore, res: List[int], from_where: List[str], select: int = 0, image_size: List[int]=[1024, 256], num_rows: int = 1, font_scale=2, thickness=4, cmap_name="plasma", save_name="null-text+ptp"): 150 | # tokens = stable.tokenizer_list[0].encode(prompts[select]) 151 | # decoder = stable.tokenizer_list[0].decode # auffusion-full 152 | tokens = stable.tokenizer.encode(prompts[select]) 153 | decoder = stable.tokenizer.decode 154 | attention_maps = aggregate_attention(prompts, attention_store, res, from_where, True, select) 155 | images = [] 156 | 157 | cmap = plt.get_cmap(cmap_name) 158 | cmap_r = cmap.reversed() 159 | 160 | for i in range(len(tokens)): 161 | image = attention_maps[:, :, i] 162 | image = 255 * image / image.max() 163 | image = image.unsqueeze(-1).expand(*image.shape, 3) 164 | image = image.numpy().astype(np.uint8) 165 | 166 | image = cmap(np.array(image)[:,:,0])[:, :, :3] 167 | image = (image - image.min()) / (image.max() - image.min()) 168 | image = Image.fromarray(np.uint8(image*255)) 169 | image = np.array(image.resize(image_size)) 170 | 171 | image = ptp_utils.text_under_image(image, decoder(int(tokens[i])), font_scale=font_scale, thickness=thickness) 172 | images.append(image) 173 | ptp_utils.view_images(np.stack(images, axis=0), num_rows=num_rows, save_name=save_name) -------------------------------------------------------------------------------- /eval/configs/audiocaps/edit_delete.csv: -------------------------------------------------------------------------------- 1 | Uid,Caption,edit_Caption 2 | 1AezQf37NWA_50000_60000,A man speaks and then whistles,A man speaks and then 3 | z_cfAqGxyT4_40000_50000,Wind blowing and water splashing,Wind blowing and 4 | X1zntKNqjYE_110000_120000,A man speaking as birds are chirping,A man speaking as 5 | TV2uGumcGIs_190000_200000,A man is speaking as paper is crumpling,A man is speaking as 6 | HIQq6nxHI_M_30000_40000,A motorboat engine running as wind blows into a microphone,A motorboat engine running as into a microphone 7 | sZsLJe0VAvA_540000_550000,A helicopter flying as wind blows into a microphone,A helicopter flying as into a microphone 8 | y9RB_DvaqcI_30000_40000,A crowd of people applauding and cheering,A crowd of people and cheering 9 | 19vv-oL7w5c_0_10000,"Water is falling, splashing and gurgling","Water is falling, splashing and" 10 | 912ulwC5d6M_25000_35000,A baby crying and breathing,A baby crying and 11 | 8e6KCFH_58g_170000_180000,A man speaking and birds chirping,A man speaking and 12 | mPf3TRcdH-E_130000_140000,Birds chirping and tweeting,Birds and tweeting 13 | ugxHt-50Z3M_530000_540000,A woman speaks as food sizzles in a pan,A woman speaks as in a pan 14 | RVUGxY7T0Lc_70000_80000,Ocean waves crashing as wind blows into a microphone,Ocean waves crashing as into a microphone 15 | f2ztuNPUUVk_15000_25000,A woman speaks and a cat meows,A woman speaks and a 16 | gD-RsNa2m-g_140000_150000,Typing on a keyboard with a man speaking,Typing on a keyboard with a 17 | YM8auY6sjAs_140000_150000,An idle vehicle engine running as wind blows into a microphone,An idle vehicle engine running as into a microphone 18 | fQbHCvgLWlY_50000_60000,A man speaks while water splashes,A man speaks while 19 | IfvNGeTUJ6U_10000_20000,A vehicle engine running and revving several times,A vehicle engine running and several times 20 | 8uS54KknUZA_30000_40000,Water trickling while a man speaks,Water trickling while 21 | Fm9MMRvkax0_50000_60000,A girl laughing as a person is snoring,A girl laughing as a person is 22 | 8LjypIOwEns_40000_50000,A powerful engine revs as it idles,A powerful engine revs as it 23 | eiPzGovKdZs_160000_170000,Water splashes and wind blows,Water splashes and 24 | XXVQ1NVI7Bk_100000_110000,Rough sanding and scraping,Rough and scraping 25 | 3oT4NTXrokU_490000_500000,A woman and a child speaking,A woman and a speaking 26 | 1ahjgFjrMbc_30000_40000,A man is speaking and a dog is barking,A man is speaking and a 27 | 2H3miNoF_Z8_30000_40000,A cat meows and hisses,A cat meows and 28 | 20JihmwTTr8_30000_40000,Race cars are passing by,Race are passing by 29 | Awre04t_FEY_15000_25000,Engine running and revving,Engine running and 30 | FmRyBPgemi4_190000_200000,A man is speaking as birds are chirping,A man is speaking as 31 | zZ3FzOY06CI_30000_40000,Gurgling and splashing water,Gurgling and 32 | fKPuVSiWpR4_560000_570000,Male speaking and birds chirping,Male speaking and 33 | f_sOWxFAi1g_30000_40000,Humming of an engine with people speaking,Humming of an engine with 34 | KU-A8g4wCU4_100000_110000,A man talks and then whistles,A man talks and then 35 | 2Sb8aPFdwFU_30000_40000,Water splashing and birds chirping,Water splashing and 36 | ylbRydhutFc_30000_40000,A man talking as pigeons coo and bird wings flap,A man talking as and bird wings flap 37 | l-l9_Byf9hg_240000_250000,A man speaks as water splashes,A man speaks as 38 | 5FrqyZr00wQ_2000_12000,A car horn honks and echoes,A car horn honks and 39 | 9OVV5Gd6pkM_14000_24000,An engine revving as it idles,An engine revving as 40 | LFVvHMNGhmI_30000_40000,Ocean waves crashing as wind blows heavily into a microphone,Ocean waves crashing as heavily into a microphone 41 | yGutSfyK4y8_30000_40000,Water splashing followed by a woman speaking,Water splashing followed by a 42 | mPf3TRcdH-E_130000_140000,Birds chirping and tweeting,Birds and tweeting 43 | B6bTIIl46fA_30000_40000,Water splashes and wind blows as men speak,Water splashes and as men speak 44 | xesCkDs3tQY_30000_40000,A vehicle engine revving and accelerating,A vehicle engine revving and 45 | BFih5FgzcRU_30000_40000,A man speaking with murmuring in the background,A man speaking with in the background 46 | 48sv95kLSVI_530000_540000,An adult male is speaking and typing on a keyboard,An adult male is speaking and 47 | 3i3V8pd6xOg_30000_40000,A woman speaks and a crowd applauds,A woman speaks and a crowd 48 | G-SUDENc2sc_30000_40000,A man and woman speak,A man and speak 49 | n67Nl29gqKk_90000_100000,A baby cries and people speak,A and people speak 50 | qcqZflfGeOw_10000_20000,A baby and a woman laugh,A baby and a woman 51 | soLjCy73csI_0_10000,An aircraft flying in the distance as wind blows into a microphone,An aircraft flying in the distance as into a microphone 52 | eiPzGovKdZs_160000_170000,Water splashes and wind blows,Water splashes and 53 | bS2PSfgW6J4_60000_70000,People laugh and speak,People and speak 54 | B4Aas8bd_Zg_30000_40000,Emergency sirens and muffled speech,and muffled speech 55 | qnIrh8pROYo_20000_30000,Water splashes and a man speaks,Water splashes and a 56 | Qdw-mzhWh-I_30000_40000,Gunshots and explosions,and explosions 57 | 8K2szL8ZtFU_30000_40000,A cat meowing and whining,A cat and whining 58 | y7_FciM9kXY_0_10000,A toilet flushes and water drains,A toilet flushes and 59 | EHBlH_kffh4_0_10000,Humming of engines with people speaking,Humming of engines with 60 | RVUGxY7T0Lc_70000_80000,Ocean waves crashing as wind blows into a microphone,Ocean waves crashing as into a microphone 61 | 3oaXDeQTEPQ_30000_40000,A man is giving a speech and a crowd cheers,A man is giving a speech and a 62 | tjq0Z9IJw7Q_250000_260000,A man speaks as water runs,A man speaks as 63 | VlusD3K-8iY_190000_200000,Birds chirping and singing,Birds and singing 64 | HH9n0qP-HCM_30000_40000,An engine idling with light wind,An engine idling with 65 | HQYDO1uOreg_150000_160000,Whistling with music playing,with music playing 66 | BPNURHcm86I_0_10000,Humming of an engine with sirens ringing,Humming of an engine with 67 | 8KwqaI9PJuk_7000_17000,A crowd applauding and cheering,A crowd and cheering 68 | jtMmwfxDEKA_30000_40000,A baby crying and cooing,A baby crying and 69 | DvJAnh2ywLI_30000_40000,A man speaks as pigeons coo and wings flap,A man speaks as pigeons coo and 70 | 779sKwaX_Hc_120000_130000,A woman speaks and food sizzles,A woman speaks and 71 | Lsv7HQizztA_0_10000,A car engine revs as it passes by,A car as it passes by 72 | DlyUeQGh1hE_110000_120000,A man speaking followed by a crowd of people applauding and cheering,A man speaking followed by a crowd of people and cheering 73 | a6Vih-pDIiI_0_10000,An engine starting and running,An engine and running 74 | 4i7Z_Z8gd7E_120000_130000,A crowd cheers and a man speaks,A crowd cheers and a 75 | y9RB_DvaqcI_30000_40000,A crowd of people applauding and cheering,A crowd of people and cheering 76 | 50Qv80Ykg7U_30000_40000,A woman and child speak,A woman and speak 77 | hUUvJlkjso8_110000_120000,Wind blowing and waves crashing,and waves crashing 78 | mheuTFdwNcE_140000_150000,A crowd of people cheering and applauding,A crowd of people and applauding 79 | G-SUDENc2sc_30000_40000,A man and woman speak,A man and speak 80 | RVUGxY7T0Lc_70000_80000,Ocean waves crashing as wind blows into a microphone,Ocean waves crashing as into a microphone 81 | jpqUA3PKge4_140000_150000,A man speaks as food sizzles,A man speaks as 82 | z8MRUCmCjyE_30000_40000,Continuous clanking and rustling,Continuous clanking and 83 | u0V7wclNZlU_30000_40000,Birds chirping and rustling,Birds and rustling 84 | fRZs8uHQCpg_10000_20000,A mid-size motor vehicle engine is idling and vibrating,A mid-size motor vehicle engine is idling and 85 | 6XwkRqJoffU_30000_40000,Emergency vehicle driving with siren on,Emergency vehicle driving with on 86 | 8sZtN1zgPwA_130000_140000,Water flowing and splashing,Water flowing and 87 | 27hWTLzVf3U_80000_90000,A vehicle passes by,A vehicle passes by 88 | VbCQjFyiR2I_170000_180000,Water flows and splashes,Water flows and 89 | zZ3FzOY06CI_30000_40000,Gurgling and splashing water,Gurgling and 90 | zD7X2XSrEFM_30000_40000,Pigeons coo and flap their wings,Pigeons and flap their wings 91 | KhIivfOIuYI_80000_90000,A man is speaking as food is frying,A man is speaking as food is 92 | yE89ARB4Veo_30000_40000,Pigeons cooing and flapping their wings,Pigeons and flapping their wings 93 | 2JdEPktKxj8_30000_40000,"Dogs are barking and growling, and an adult male speaks",", and an adult male speaks" 94 | 7jx1bu0M9fk_140000_150000,Water splashes and gurgles,Water splashes and 95 | 3Vs6jCyH5vg_10000_20000,Snoring and heavy breathing,Snoring and heavy 96 | ylbRydhutFc_30000_40000,A man talking as pigeons coo and bird wings flap,A as pigeons coo and bird wings flap 97 | nQI-AxPN3vc_30000_40000,Humming of an engine with some rustling,Humming of an engine with some 98 | 4bvp9YmvRmU_30000_40000,People are talking and water is splashing,People are talking and 99 | EcDQzSQ-d_I_30000_40000,A baby crying and a man speaking,A baby crying and 100 | -------------------------------------------------------------------------------- /suppresseot/run_and_display.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from tqdm import tqdm 4 | from typing import Optional, List 5 | 6 | from prompt2prompt import ptp_utils 7 | from prompt2prompt.attn_control import aggregate_attention 8 | from suppresseot import wo_utils 9 | from suppresseot.attn_loss import AttnLoss 10 | 11 | 12 | LOW_RESOURCE = True 13 | 14 | def update_context(context: torch.Tensor, loss: torch.Tensor, scale: int, factor: float) -> torch.Tensor: 15 | """ 16 | Update the text embeddings according to the attention loss. 17 | 18 | :param context: text embeddings to be updated 19 | :param loss: ours loss 20 | :param factor: factor for update text embeddings. 21 | :return: 22 | """ 23 | grad_cond = torch.autograd.grad(outputs=loss.requires_grad_(True), inputs=[context], retain_graph=False)[0] 24 | context = context - (scale * factor) * grad_cond 25 | return context 26 | 27 | @torch.no_grad() 28 | def text2image_ldm_stable( 29 | model, 30 | prompt: List[str], 31 | controller, 32 | num_inference_steps: int = 50, 33 | guidance_scale: Optional[float] = 7.5, 34 | generator: Optional[torch.Generator] = None, 35 | latent: Optional[torch.FloatTensor] = None, 36 | uncond_embeddings=None, 37 | start_time=100, 38 | return_type='image' 39 | ): 40 | ptp_utils.register_attention_control(model, controller) 41 | height = 256 42 | width = 1024 43 | batch_size = len(prompt) 44 | 45 | # text_input = model.tokenizer_list[0]( # auffusion-full 46 | text_input = model.tokenizer( 47 | prompt, 48 | padding="max_length", 49 | # max_length=model.tokenizer_list[0].model_max_length, # auffusion-full 50 | max_length=model.tokenizer.model_max_length, 51 | truncation=True, 52 | return_tensors="pt", 53 | ) 54 | # text_embeddings = model.text_encoder_list[0](text_input.input_ids.to(model.device))[0] # auffusion-full 55 | text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0] 56 | 57 | max_length = text_input.input_ids.shape[-1] 58 | if uncond_embeddings is None: 59 | # uncond_input = model.tokenizer_list[0]([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt") # auffusion-full 60 | uncond_input = model.tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt") 61 | # uncond_embeddings_ = model.text_encoder_list[0](uncond_input.input_ids.to(model.device))[0] # auffusion-full 62 | uncond_embeddings_ = model.text_encoder(uncond_input.input_ids.to(model.device))[0] 63 | scale = 20 64 | else: 65 | uncond_embeddings_ = None 66 | scale = 5 67 | 68 | latent, _ = ptp_utils.init_latent(latent, model, height, width, generator, batch_size) 69 | 70 | _latent, _latent_erase = latent.clone().to(model.device), latent.clone().to(model.device) 71 | latents = torch.cat([_latent, _latent_erase]) 72 | 73 | attn_loss_func = AttnLoss(model.device, 'cosine', controller.n, controller.token_indices, 74 | controller.lambda_retain, controller.lambda_erase, controller.lambda_self_retain, controller.lambda_self_erase) 75 | 76 | model.scheduler.set_timesteps(num_inference_steps) 77 | # text embedding for erasing 78 | controller.text_embeddings_erase = text_embeddings.clone() 79 | 80 | scale_range = np.linspace(1., .1, len(model.scheduler.timesteps)) 81 | pbar = tqdm(model.scheduler.timesteps[-start_time:], desc='Suppress EOT', ncols=100, colour="red") 82 | for i, t in enumerate(pbar): 83 | if uncond_embeddings_ is None: 84 | context = torch.cat([uncond_embeddings[i].expand(*text_embeddings.shape), text_embeddings]) 85 | if LOW_RESOURCE: 86 | context = (uncond_embeddings[i].expand(*text_embeddings.shape), text_embeddings) 87 | else: 88 | context = torch.cat([uncond_embeddings_, text_embeddings]) 89 | if LOW_RESOURCE: 90 | context = (uncond_embeddings_, text_embeddings) 91 | controller.i = i 92 | 93 | # conditional branch: erase content for text embeddings 94 | if controller.i >= controller.cross_retain_steps: 95 | controller.text_embeddings_erase = \ 96 | wo_utils.woword_eot_context(text_embeddings.clone(), controller.token_indices, controller.alpha, 97 | controller.method, controller.n) 98 | 99 | controller.baseline = True 100 | if controller.MAX_STEP_TO_ERASE > controller.i >= controller.cross_retain_steps and not (controller.text_embeddings_erase == text_embeddings).all() and \ 101 | (attn_loss_func.lambda_retain or attn_loss_func.lambda_erase or attn_loss_func.lambda_self_retain or attn_loss_func.lambda_self_erase): 102 | controller.uncond = False 103 | controller.cur_att_layer = 32 # w=1, skip unconditional branch 104 | controller.attention_store = {} 105 | noise_prediction_text = model.unet(_latent, t, encoder_hidden_states=text_embeddings)["sample"] 106 | attention_maps = aggregate_attention(controller, 16, ["up", "down"], is_cross=True) 107 | self_attention_maps = aggregate_attention(controller, 16, ["up", "down", "mid"], is_cross=False) 108 | 109 | del noise_prediction_text 110 | # update controller.text_embeddings_erase for some timestep 111 | iter = controller.iter_each_step 112 | while iter > 0: 113 | with torch.enable_grad(): 114 | controller.cur_att_layer = 32 # w=1, skip unconditional branch 115 | controller.attention_store = {} 116 | # conditional branch 117 | text_embeddings_erase = controller.text_embeddings_erase.clone().detach().requires_grad_(True) 118 | # forward pass of conditional branch with text_embeddings_erase 119 | noise_prediction_text = model.unet(_latent_erase, t, encoder_hidden_states=text_embeddings_erase)["sample"] 120 | model.unet.zero_grad() 121 | attention_maps_erase = aggregate_attention(controller, 16, ["up", "down", "mid"], is_cross=True) 122 | self_attention_maps_erase = aggregate_attention(controller, 16, ["up", "down", "mid"], is_cross=False) 123 | 124 | # attention loss 125 | loss = attn_loss_func(attention_maps, attention_maps_erase, self_attention_maps, self_attention_maps_erase) 126 | if loss != .0: 127 | pbar.set_postfix({'loss': loss if isinstance(loss, float) else loss.item()}) 128 | text_embeddings_erase = update_context(context=text_embeddings_erase, loss=loss, 129 | scale=scale, factor=np.sqrt(scale_range[i])) 130 | del noise_prediction_text 131 | torch.cuda.empty_cache() 132 | controller.text_embeddings_erase = text_embeddings_erase.clone().detach().requires_grad_(False) 133 | iter -= 1 134 | 135 | # "uncond_embeddings_ is None" for real images, "uncond_embeddings_ is not None" for generated images. 136 | context_erase = (uncond_embeddings[i].expand(*text_embeddings.shape), controller.text_embeddings_erase) \ 137 | if uncond_embeddings_ is None else (uncond_embeddings_, controller.text_embeddings_erase) 138 | controller.attention_store = {} 139 | controller.baseline = False 140 | contexts = [torch.cat([context[0], context_erase[0]]), torch.cat([context[1], context_erase[1]])] 141 | latents = ptp_utils.diffusion_step(model, controller, latents, contexts, t, guidance_scale, low_resource=LOW_RESOURCE) 142 | _latent, _latent_erase = latents 143 | _latent, _latent_erase = _latent.unsqueeze(0), _latent_erase.unsqueeze(0) 144 | 145 | if return_type == 'image': 146 | latents = 1 / 0.18215 * latents.detach() 147 | _image = model.vae.decode(latents)['sample'] 148 | _image = (_image / 2 + 0.5).clamp(0, 1) 149 | null_edit = _image.cpu().permute(0, 2, 3, 1).numpy()[0] 150 | eot_edit = _image.cpu().permute(0, 2, 3, 1).numpy()[1] 151 | image = np.stack((null_edit, eot_edit), axis=0) 152 | else: 153 | image = latents 154 | return image, latent 155 | 156 | def run_and_display(ldm_stable, prompts, controller, latent=None, generator=None, uncond_embeddings=None, args=None): 157 | images, x_t = text2image_ldm_stable(ldm_stable, prompts, controller, latent=latent, 158 | num_inference_steps=args.num_inference_steps, guidance_scale=args.guidance_scale, 159 | generator=generator, uncond_embeddings=uncond_embeddings) 160 | return images, x_t -------------------------------------------------------------------------------- /null_text_inversion/null_text_inversion.py: -------------------------------------------------------------------------------- 1 | # https://github.com/google/prompt-to-prompt 2 | 3 | from typing import Optional, Union, Tuple, List, Callable, Dict 4 | from tqdm import tqdm 5 | import torch 6 | 7 | from diffusers import StableDiffusionPipeline, DDIMScheduler 8 | import torch.nn.functional as nnf 9 | import numpy as np 10 | # import seq_aligner 11 | from torch.optim.adam import Adam 12 | from PIL import Image 13 | 14 | # from prompt2prompt import ptp_utils 15 | 16 | NUM_DDIM_STEPS = 100 17 | GUIDANCE_SCALE_TRAIN = 7.5 18 | 19 | class NullInversion: 20 | 21 | def prev_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, 22 | sample: Union[torch.FloatTensor, np.ndarray]): 23 | prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps 24 | alpha_prod_t = self.scheduler.alphas_cumprod[timestep] 25 | alpha_prod_t_prev = self.scheduler.alphas_cumprod[ 26 | prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod 27 | beta_prod_t = 1 - alpha_prod_t 28 | pred_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5 29 | pred_sample_direction = (1 - alpha_prod_t_prev) ** 0.5 * model_output 30 | prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction 31 | return prev_sample 32 | 33 | def next_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, 34 | sample: Union[torch.FloatTensor, np.ndarray]): 35 | timestep, next_timestep = min( 36 | timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps, 999), timestep 37 | alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod 38 | alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep] 39 | beta_prod_t = 1 - alpha_prod_t 40 | next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5 41 | next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output 42 | next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction 43 | return next_sample 44 | 45 | def get_noise_pred_single(self, latents, t, context): 46 | noise_pred = self.model.unet(latents, t, encoder_hidden_states=context)["sample"] 47 | return noise_pred 48 | 49 | def get_noise_pred(self, latents, t, is_forward=True, context=None): 50 | latents_input = torch.cat([latents] * 2) 51 | if context is None: 52 | context = self.context 53 | guidance_scale = 1 if is_forward else GUIDANCE_SCALE_TRAIN 54 | noise_pred = self.model.unet(latents_input, t, encoder_hidden_states=context)["sample"] 55 | noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2) 56 | noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) 57 | if is_forward: 58 | latents = self.next_step(noise_pred, t, latents) 59 | else: 60 | latents = self.prev_step(noise_pred, t, latents) 61 | return latents 62 | 63 | @torch.no_grad() 64 | def latent2image(self, latents, return_type='np'): 65 | latents = 1 / 0.18215 * latents.detach() 66 | image = self.model.vae.decode(latents)['sample'] 67 | if return_type == 'np': 68 | image = (image / 2 + 0.5).clamp(0, 1) 69 | image = image.cpu().permute(0, 2, 3, 1).numpy()[0] 70 | return image 71 | 72 | @torch.no_grad() 73 | def image2latent(self, image): 74 | with torch.no_grad(): 75 | if type(image) is Image: 76 | image = np.array(image) 77 | if type(image) is torch.Tensor and image.dim() == 4: 78 | latents = image 79 | else: 80 | image = torch.from_numpy(image).float() / 0.5 - 1 81 | image = image.permute(2, 0, 1).unsqueeze(0).to(self.model.device) 82 | latents = self.model.vae.encode(image)['latent_dist'].mean 83 | latents = latents * 0.18215 84 | return latents 85 | 86 | @torch.no_grad() 87 | def init_prompt(self, prompt: str): 88 | uncond_input = self.tokenizer( 89 | [""], padding="max_length", max_length=self.tokenizer.model_max_length, 90 | return_tensors="pt" 91 | ) 92 | # uncond_embeddings = self.model.text_encoder_list[0](uncond_input.input_ids.to(self.model.device))[0] # auffusion-full 93 | uncond_embeddings = self.model.text_encoder(uncond_input.input_ids.to(self.model.device))[0] 94 | text_input = self.tokenizer( 95 | [prompt], 96 | padding="max_length", 97 | max_length=self.tokenizer.model_max_length, 98 | truncation=True, 99 | return_tensors="pt", 100 | ) 101 | # text_embeddings = self.model.text_encoder_list[0](text_input.input_ids.to(self.model.device))[0] # auffusion-full 102 | text_embeddings = self.model.text_encoder(text_input.input_ids.to(self.model.device))[0] 103 | self.context = torch.cat([uncond_embeddings, text_embeddings]) 104 | self.prompt = prompt 105 | 106 | @torch.no_grad() 107 | def ddim_loop(self, latent): 108 | uncond_embeddings, cond_embeddings = self.context.chunk(2) 109 | all_latent = [latent] 110 | latent = latent.clone().detach() 111 | for i in range(NUM_DDIM_STEPS): 112 | t = self.model.scheduler.timesteps[len(self.model.scheduler.timesteps) - i - 1] 113 | noise_pred = self.get_noise_pred_single(latent, t, cond_embeddings) 114 | latent = self.next_step(noise_pred, t, latent) 115 | all_latent.append(latent) 116 | return all_latent 117 | 118 | @property 119 | def scheduler(self): 120 | return self.model.scheduler 121 | 122 | @torch.no_grad() 123 | def ddim_inversion(self, image): 124 | latent = self.image2latent(image) 125 | image_rec = self.latent2image(latent) 126 | ddim_latents = self.ddim_loop(latent) 127 | return image_rec, ddim_latents 128 | 129 | def null_optimization(self, latents, num_inner_steps, epsilon): 130 | uncond_embeddings, cond_embeddings = self.context.chunk(2) 131 | uncond_embeddings_list = [] 132 | latent_cur = latents[-1] 133 | bar = tqdm(total=num_inner_steps * NUM_DDIM_STEPS) 134 | for i in range(NUM_DDIM_STEPS): 135 | uncond_embeddings = uncond_embeddings.clone().detach() 136 | uncond_embeddings.requires_grad = True 137 | optimizer = Adam([uncond_embeddings], lr=1e-4 * (1. - i / 1000.)) 138 | latent_prev = latents[len(latents) - i - 2] 139 | t = self.model.scheduler.timesteps[i] 140 | with torch.no_grad(): 141 | noise_pred_cond = self.get_noise_pred_single(latent_cur, t, cond_embeddings) 142 | for j in range(num_inner_steps): 143 | noise_pred_uncond = self.get_noise_pred_single(latent_cur, t, uncond_embeddings) 144 | noise_pred = noise_pred_uncond + GUIDANCE_SCALE_TRAIN * (noise_pred_cond - noise_pred_uncond) 145 | latents_prev_rec = self.prev_step(noise_pred, t, latent_cur) 146 | loss = nnf.mse_loss(latents_prev_rec, latent_prev) 147 | optimizer.zero_grad() 148 | loss.backward() 149 | optimizer.step() 150 | loss_item = loss.item() 151 | bar.update() 152 | if loss_item < epsilon + i * 2e-5: 153 | break 154 | for j in range(j + 1, num_inner_steps): 155 | bar.update() 156 | uncond_embeddings_list.append(uncond_embeddings[:1].detach()) 157 | with torch.no_grad(): 158 | context = torch.cat([uncond_embeddings, cond_embeddings]) 159 | latent_cur = self.get_noise_pred(latent_cur, t, False, context) 160 | bar.close() 161 | return uncond_embeddings_list 162 | 163 | def invert(self, image_path: str, prompt: str, offsets=(0, 0, 0, 0), num_inner_steps=10, early_stop_epsilon=1e-5, 164 | inversion='Null-text', verbose=False): 165 | self.init_prompt(prompt) 166 | image_gt = image_path 167 | if verbose: 168 | print("DDIM inversion...") 169 | image_rec, ddim_latents = self.ddim_inversion(image_gt) 170 | 171 | assert inversion in ['NT', 'NPI'] 172 | if inversion == 'NT': 173 | print("Null-text optimization...") 174 | uncond_embeddings = self.null_optimization(ddim_latents, num_inner_steps, early_stop_epsilon) 175 | else: 176 | print("Negative prompt inversion...") 177 | uncond_embeddings, cond_embeddings = self.context.chunk(2) 178 | uncond_embeddings = [cond_embeddings] * NUM_DDIM_STEPS 179 | return (image_gt, image_rec), ddim_latents[-1], uncond_embeddings 180 | 181 | def __init__(self, model): 182 | scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, 183 | set_alpha_to_one=False) 184 | self.model = model 185 | # self.tokenizer = self.model.tokenizer_list[0] # auffusion-full 186 | self.tokenizer = self.model.tokenizer 187 | self.model.scheduler.set_timesteps(NUM_DDIM_STEPS) 188 | self.prompt = None 189 | self.context = None -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import ast 3 | import torch 4 | import argparse 5 | import numpy as np 6 | from scipy.io.wavfile import write 7 | from diffusers.utils.import_utils import is_xformers_available 8 | 9 | from diffusers import StableDiffusionPipeline, DDIMScheduler 10 | from auffusion.auffusion_pipeline import AuffusionPipeline 11 | 12 | from null_text_inversion.null_text_inversion import NullInversion 13 | from utils.converter import load_wav, mel_spectrogram, normalize_spectrogram, denormalize_spectrogram, Generator, get_mel_spectrogram_from_audio 14 | from utils.utils import pad_spec, image_add_color, torch_to_pil, normalize, denormalize 15 | 16 | 17 | from prompt2prompt.attn_control import AttentionStore, show_cross_attention 18 | from suppresseot.run_and_display import run_and_display 19 | 20 | 21 | def parse_args(): 22 | parser = argparse.ArgumentParser() 23 | parser.add_argument("--sd_version", type=str, default='auffusion-full-no-adapter', help='version of stable diffusion model.') 24 | parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible inference.") 25 | parser.add_argument("--output_dir", type=str, default="./audio_examples/output_audios", help="The output directory where the model predictions will be written.") 26 | parser.add_argument('--prompt', type=str, default='After a gunshot, there was a burst of dog barking', help='prompt for generated or real audio') 27 | parser.add_argument('--audio_path', type=str, default='audio_examples/input_audios/After a gunshot, there was a burst of dog barking.wav', help='audio path') 28 | 29 | parser.add_argument("--enable_xformers", default=True, action="store_true", help="Whether or not to use xformers.") 30 | parser.add_argument("--guidance_scale", type=float, default=7.5, help="The scale of guidance.") 31 | parser.add_argument("--num_inference_steps", type=int, default=100, help="Number of inference steps to perform.") 32 | 33 | parser.add_argument("--width", type=int, default=1024, help="Width of the spec.") 34 | parser.add_argument("--height", type=int, default=256, help="Height of the spec.") 35 | parser.add_argument("--sample_rate", type=int, default=16000, help="The sample rate of audio.") 36 | parser.add_argument("--duration", type=int, default=10, help="The duration(s) of audio.") 37 | 38 | 39 | parser.add_argument("--input_dir", type=str, default="./audio_examples/input_audios", help="The input directory where put the audios to be edited.") 40 | parser.add_argument('--inversion', type=str, default='NT', help='NT (Null-text), NPI (Negative-prompt-inversion).') 41 | 42 | parser.add_argument('--token_indices', type=ast.literal_eval, default='[[1,]]', help='index of without words.') 43 | parser.add_argument('--cross_retain_steps', type=ast.literal_eval, default='[.2,]', help='perform the "wo" punish when step >= cross_wo_steps') 44 | parser.add_argument('--alpha', type=ast.literal_eval, default='[1.,]', help="punishment ratio") 45 | parser.add_argument('--iter_each_step', type=int, default=5, help="the number of iteration for each step to update text embedding") 46 | parser.add_argument('--max_step_to_erase', type=int, default=20, help='erase/suppress max step of diffusion model') 47 | parser.add_argument('--method', type=str, default='soft-weight', help='soft-weight, alpha, beta, delete, weight') 48 | 49 | parser.add_argument('--lambda_retain', type=float, default=1., help='lambda for cross attention retain loss') 50 | parser.add_argument('--lambda_erase', type=float, default=-.5, help='lambda for cross attention erase loss') 51 | parser.add_argument('--lambda_self_retain', type=float, default=1., help='lambda for self attention retain loss') 52 | parser.add_argument('--lambda_self_erase', type=float, default=-.5, help='lambda for self attention erase loss') 53 | 54 | args = parser.parse_args() 55 | return args 56 | 57 | def load_model(sd_version, device): 58 | if sd_version == "auffusion-full-no-adapter": 59 | pretrained_model_name_or_path = "ckpt/auffusion-full-no-adapter" 60 | ldm_stable = StableDiffusionPipeline.from_pretrained(pretrained_model_name_or_path).to(device) 61 | elif sd_version == sd_version == "sd_1_4": 62 | pretrained_model_name_or_path = "/home/jiayuhang/.cache/huggingface/hub/models--CompVis--stable-diffusion-v1-4/snapshots/133a221b8aa7292a167afc5127cb63fb5005638b" 63 | scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False, steps_offset=1) 64 | ldm_stable = StableDiffusionPipeline.from_pretrained(pretrained_model_name_or_path, scheduler=scheduler).to(device) 65 | else: 66 | raise ValueError('Unsupported stable diffusion version') 67 | return ldm_stable 68 | 69 | 70 | def load_1024_256(audio_path): 71 | audio, sampling_rate = load_wav(audio_path) 72 | audio, spec = get_mel_spectrogram_from_audio(audio) 73 | norm_spec = normalize_spectrogram(spec) 74 | norm_spec = pad_spec(norm_spec, 1024) 75 | norm_spec = norm_spec.permute(1, 2, 0).cpu().numpy() # (256, 1024, 3) # reshape 76 | return norm_spec 77 | 78 | 79 | def store_1024_256(output_spec, sample_rate, audio_path): 80 | ### vocoder 81 | device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') 82 | vocoder = Generator.from_pretrained("ckpt/auffusion-full-no-adapter", subfolder="vocoder") 83 | vocoder = vocoder.to(device=device, dtype=torch.float16) 84 | 85 | norm_spec = output_spec[:, :1000, :] 86 | norm_spec = torch.from_numpy(norm_spec).permute(2, 0, 1).to(device=device) # reshape torch.Size([3, 256, 1000]) 87 | denorm_spec = denormalize_spectrogram(norm_spec) 88 | with torch.autocast("cuda"): 89 | denorm_spec_audio = vocoder.inference(denorm_spec) 90 | write(audio_path, sample_rate, denorm_spec_audio.squeeze()) 91 | print(f"Successfully save {audio_path}") 92 | 93 | 94 | def show_spec(output_spec, audio_path): 95 | device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') 96 | norm_spec = torch.from_numpy(output_spec).permute(2, 0, 1).to(device=device) 97 | raw_image = image_add_color(torch_to_pil(norm_spec[:,:,:1000])) 98 | raw_image.save(audio_path) 99 | 100 | 101 | def edit_audio(): 102 | args = parse_args() 103 | device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') 104 | ldm_stable = load_model(args.sd_version, device) 105 | ldm_stable.set_progress_bar_config(disable=True) 106 | if is_xformers_available() and args.enable_xformers: 107 | ldm_stable.enable_xformers_memory_efficient_attention() 108 | 109 | orig_prompt = args.prompt 110 | input_audio_path = args.audio_path 111 | input_audio = load_1024_256(input_audio_path) 112 | os.makedirs(os.path.join(args.output_dir, orig_prompt), exist_ok=True) 113 | 114 | # Null-Text Inversion 115 | null_inversion = NullInversion(ldm_stable) 116 | with torch.autocast("cuda"): 117 | (audio_gt, audio_rec), x_t, uncond_embeddings = null_inversion.invert(input_audio, orig_prompt, inversion=args.inversion, \ 118 | num_inner_steps=args.iter_each_step, verbose=True) 119 | 120 | show_spec(audio_gt, f"audio_examples/output_audios/{orig_prompt}/audio_gt.png") 121 | show_spec(audio_rec, f"audio_examples/output_audios/{orig_prompt}/audio_rec.png") 122 | store_1024_256(audio_gt, args.sample_rate, f"audio_examples/output_audios/{orig_prompt}/audio_gt.wav") 123 | store_1024_256(audio_rec, args.sample_rate, f"audio_examples/output_audios/{orig_prompt}/audio_rec.wav") 124 | 125 | #SuppressEOT 126 | for token_indices in args.token_indices: 127 | for cross_retain_steps in args.cross_retain_steps: 128 | for alpha in args.alpha: 129 | controller = AttentionStore(token_indices, \ 130 | alpha, \ 131 | args.method, \ 132 | cross_retain_steps, \ 133 | len(ldm_stable.tokenizer.encode((orig_prompt))), \ 134 | args.iter_each_step, \ 135 | args.max_step_to_erase, \ 136 | lambda_retain=args.lambda_retain, \ 137 | lambda_erase=args.lambda_erase, \ 138 | lambda_self_retain=args.lambda_self_retain, \ 139 | lambda_self_erase=args.lambda_self_erase) 140 | 141 | with torch.autocast("cuda"): 142 | audio_edit, x_t = run_and_display(ldm_stable, \ 143 | [orig_prompt], \ 144 | controller, \ 145 | latent=x_t, \ 146 | uncond_embeddings=uncond_embeddings, \ 147 | args=args) 148 | 149 | 150 | show_spec(audio_edit[1], f"audio_examples/output_audios/{orig_prompt}/audio_eot_edit.png") 151 | store_1024_256(audio_edit[1], args.sample_rate, f"audio_examples/output_audios/{orig_prompt}/audio_eot_edit.wav") 152 | 153 | 154 | def gen_audio(): 155 | args = parse_args() 156 | device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') 157 | ldm_stable = load_model(args.sd_version, device) 158 | ldm_stable.set_progress_bar_config(disable=True) 159 | if is_xformers_available() and args.enable_xformers: 160 | ldm_stable.enable_xformers_memory_efficient_attention() 161 | 162 | generator = torch.Generator(device=device).manual_seed(args.seed) 163 | text_prompt = args.edit_prompt 164 | audio_path = os.path.join(args.input_dir, text_prompt + '.wav') 165 | 166 | with torch.autocast("cuda"): 167 | output = ldm_stable(prompt = text_prompt, \ 168 | num_inference_steps = args.num_inference_steps, \ 169 | guidance_scale = args.guidance_scale, \ 170 | generator = generator, \ 171 | width = args.width, \ 172 | height = args.height, 173 | output_type="pt") 174 | 175 | if args.sd_version == "auffusion": 176 | audio_length = args.sample_rate * args.duration 177 | audio = output.audios[0][:audio_length] 178 | write(audio_path, args.sample_rate, audio) 179 | elif args.sd_version == "auffusion-full-no-adapter": 180 | output_spec = output.images[0].permute(1, 2, 0).cpu().numpy() 181 | store_1024_256(output_spec, args.sample_rate, audio_path) 182 | else: 183 | raise ValueError('Unsupported stable diffusion version') 184 | 185 | 186 | def main(): 187 | # gen_audio() 188 | edit_audio() 189 | 190 | 191 | if __name__=="__main__": 192 | main() -------------------------------------------------------------------------------- /utils/utils.py: -------------------------------------------------------------------------------- 1 | import os, json 2 | import math, random 3 | from multiprocessing import Pool 4 | from tqdm import tqdm 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from PIL import Image 10 | import matplotlib.pyplot as plt 11 | from torchvision import transforms 12 | from transformers import CLIPTextModel 13 | from transformers import PretrainedConfig 14 | 15 | 16 | def pad_spec(spec, spec_length, pad_value=0, random_crop=True): # spec: [3, mel_dim, spec_len] 17 | assert spec_length % 8 == 0, "spec_length must be divisible by 8" 18 | if spec.shape[-1] < spec_length: 19 | # pad spec to spec_length 20 | spec = F.pad(spec, (0, spec_length - spec.shape[-1]), value=pad_value) 21 | else: 22 | # random crop 23 | if random_crop: 24 | start = random.randint(0, spec.shape[-1] - spec_length) 25 | spec = spec[:, :, start:start+spec_length] 26 | else: 27 | spec = spec[:, :, :spec_length] 28 | return spec 29 | 30 | 31 | def load_spec(spec_path): 32 | if spec_path.endswith(".pt"): 33 | spec = torch.load(spec_path, map_location="cpu") 34 | elif spec_path.endswith(".npy"): 35 | spec = torch.from_numpy(np.load(spec_path)) 36 | else: 37 | raise ValueError(f"Unknown spec file type {spec_path}") 38 | assert len(spec.shape) == 3, f"spec shape must be [3, mel_dim, spec_len], got {spec.shape}" 39 | if spec.size(0) == 1: 40 | spec = spec.repeat(3, 1, 1) 41 | return spec 42 | 43 | 44 | def random_crop_spec(spec, target_spec_length, pad_value=0, frame_per_sec=100, time_step=5): # spec: [3, mel_dim, spec_len] 45 | assert target_spec_length % 8 == 0, "spec_length must be divisible by 8" 46 | 47 | spec_length = spec.shape[-1] 48 | full_s = math.ceil(spec_length / frame_per_sec / time_step) * time_step # get full seconds(ceil) 49 | start_s = random.randint(0, math.floor(spec_length / frame_per_sec / time_step)) * time_step # random get start seconds 50 | 51 | end_s = min(start_s + math.ceil(target_spec_length / frame_per_sec), full_s) # get end seconds 52 | 53 | spec = spec[:, :, start_s * frame_per_sec : end_s * frame_per_sec] # get spec in seconds(crop more than target_spec_length because ceiling) 54 | 55 | if spec.shape[-1] < target_spec_length: 56 | spec = F.pad(spec, (0, target_spec_length - spec.shape[-1]), value=pad_value) # pad to target_spec_length 57 | else: 58 | spec = spec[:, :, :target_spec_length] # crop to target_spec_length 59 | 60 | return spec, int(start_s), int(end_s), int(full_s) 61 | 62 | 63 | 64 | def load_condion_embed(text_embed_path): 65 | if text_embed_path.endswith(".pt"): 66 | text_embed_list = torch.load(text_embed_path, map_location="cpu") 67 | elif text_embed_path.endswith(".npy"): 68 | text_embed_list = torch.from_numpy(np.load(text_embed_path)) 69 | else: 70 | raise ValueError(f"Unknown text embedding file type {text_embed_path}") 71 | if type(text_embed_list) == list: 72 | text_embed = random.choice(text_embed_list) 73 | if len(text_embed.shape) == 3: # [1, text_len, text_dim] 74 | text_embed = text_embed.squeeze(0) # random choice and return text_emb: [text_len, text_dim] 75 | return text_embed.detach().cpu() 76 | 77 | 78 | def process_condition_embed(cond_emb, max_length): # [text_len, text_dim], Padding 0 and random drop by CFG 79 | if cond_emb.shape[0] < max_length: 80 | cond_emb = F.pad(cond_emb, (0, 0, 0, max_length - cond_emb.shape[0]), value=0) 81 | else: 82 | cond_emb = cond_emb[:max_length, :] 83 | return cond_emb 84 | 85 | 86 | def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str): 87 | text_encoder_config = PretrainedConfig.from_pretrained( 88 | pretrained_model_name_or_path 89 | ) 90 | model_class = text_encoder_config.architectures[0] 91 | 92 | if model_class == "CLIPTextModel": 93 | from transformers import CLIPTextModel 94 | return CLIPTextModel 95 | if "t5" in model_class.lower(): 96 | from transformers import T5EncoderModel 97 | return T5EncoderModel 98 | if "clap" in model_class.lower(): 99 | from transformers import ClapTextModelWithProjection 100 | return ClapTextModelWithProjection 101 | else: 102 | raise ValueError(f"{model_class} is not supported.") 103 | 104 | 105 | 106 | def str2bool(string): 107 | str2val = {"True": True, "False": False, "true": True, "false": False, "none": False, "None": False} 108 | if string in str2val: 109 | return str2val[string] 110 | else: 111 | raise ValueError(f"Expected one of {set(str2val.keys())}, got {string}") 112 | 113 | 114 | def str2str(string): 115 | if string.lower() == "none" or string.lower() == "null" or string.lower() == "false" or string == "": 116 | return None 117 | else: 118 | return string 119 | 120 | 121 | def json_dump(data_json, json_save_path): 122 | with open(json_save_path, 'w') as f: 123 | json.dump(data_json, f, indent=4) 124 | f.close() 125 | 126 | 127 | def json_load(json_path): 128 | with open(json_path, 'r') as f: 129 | data = json.load(f) 130 | f.close() 131 | return data 132 | 133 | 134 | def load_json_list(path): 135 | with open(path, 'r', encoding='utf-8') as f: 136 | return [json.loads(line) for line in f.readlines()] 137 | 138 | 139 | def save_json_list(data, path): 140 | with open(path, 'w', encoding='utf-8') as f: 141 | for d in data: 142 | f.write(json.dumps(d) + '\n') 143 | 144 | 145 | def multiprocess_function(func, func_args, n_jobs=32): 146 | with Pool(processes=n_jobs) as p: 147 | with tqdm(total=len(func_args)) as pbar: 148 | for i, _ in enumerate(p.imap_unordered(func, func_args)): 149 | pbar.update() 150 | 151 | 152 | def image_add_color(spec_img): 153 | cmap = plt.get_cmap('viridis') 154 | cmap_r = cmap.reversed() 155 | image = cmap(np.array(spec_img)[:,:,0])[:, :, :3] 156 | image = (image - image.min()) / (image.max() - image.min()) 157 | image = Image.fromarray(np.uint8(image*255)) 158 | return image 159 | 160 | 161 | @staticmethod 162 | def pt_to_numpy(images: torch.FloatTensor) -> np.ndarray: 163 | """ 164 | Convert a PyTorch tensor to a NumPy image. 165 | """ 166 | images = images.cpu().permute(0, 2, 3, 1).float().numpy() 167 | return images 168 | 169 | 170 | def numpy_to_pil(images): 171 | """ 172 | Convert a numpy image or a batch of images to a PIL image. 173 | """ 174 | if images.ndim == 3: 175 | images = images[None, ...] 176 | images = (images * 255).round().astype("uint8") 177 | if images.shape[-1] == 1: 178 | # special case for grayscale (single channel) images 179 | pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] 180 | else: 181 | pil_images = [Image.fromarray(image) for image in images] 182 | 183 | return pil_images 184 | 185 | ### CODE FOR INPAITING ### 186 | def normalize(images): 187 | """ 188 | Normalize an image array to [-1,1]. 189 | """ 190 | if images.min() >= 0: 191 | return 2.0 * images - 1.0 192 | else: 193 | return images 194 | 195 | def denormalize(images): 196 | """ 197 | Denormalize an image array to [0,1]. 198 | """ 199 | if images.min() < 0: 200 | return (images / 2 + 0.5).clamp(0, 1) 201 | else: 202 | return images.clamp(0, 1) 203 | 204 | 205 | def prepare_mask_and_masked_image(image, mask): 206 | """ 207 | Prepare a binary mask and the masked image. 208 | 209 | Parameters: 210 | - image (torch.Tensor): The input image tensor of shape [3, height, width] with values in the range [0, 1]. 211 | - mask (torch.Tensor): The input mask tensor of shape [1, height, width]. 212 | 213 | Returns: 214 | - tuple: A tuple containing the binary mask and the masked image. 215 | """ 216 | # Noralize image to [0,1] 217 | if image.max() > 1: 218 | image = (image - image.min()) / (image.max() - image.min()) 219 | # Normalize image from [0,1] to [-1,1] 220 | if image.min() >= 0: 221 | image = normalize(image) 222 | # Apply the mask to the image 223 | masked_image = image * (mask < 0.5) 224 | 225 | return mask, masked_image 226 | 227 | 228 | def torch_to_pil(image): 229 | """ 230 | Convert a torch tensor to a PIL image. 231 | """ 232 | if image.min() < 0: 233 | image = denormalize(image) 234 | 235 | return transforms.ToPILImage()(image.cpu().detach().squeeze()) 236 | 237 | 238 | 239 | # class TextEncoderAdapter(nn.Module): 240 | # def __init__(self, hidden_size, cross_attention_dim=768): 241 | # super(TextEncoderAdapter, self).__init__() 242 | # self.hidden_size = hidden_size 243 | # self.cross_attention_dim = cross_attention_dim 244 | # self.proj = nn.Linear(self.hidden_size, self.cross_attention_dim) 245 | # self.norm = torch.nn.LayerNorm(self.cross_attention_dim) 246 | 247 | # def forward(self, x): 248 | # x = self.proj(x) 249 | # x = self.norm(x) 250 | # return x 251 | 252 | # def save_pretrained(self, save_directory, subfolder=""): 253 | # if subfolder: 254 | # save_directory = os.path.join(save_directory, subfolder) 255 | # os.makedirs(save_directory, exist_ok=True) 256 | # ckpt_path = os.path.join(save_directory, "adapter.pt") 257 | # config_path = os.path.join(save_directory, "config.json") 258 | # config = {"hidden_size": self.hidden_size, "cross_attention_dim": self.cross_attention_dim} 259 | # json_dump(config, config_path) 260 | # torch.save(self.state_dict(), ckpt_path) 261 | # print(f"Saving adapter model to {ckpt_path}") 262 | 263 | # @classmethod 264 | # def from_pretrained(cls, load_directory, subfolder=""): 265 | # if subfolder: 266 | # load_directory = os.path.join(load_directory, subfolder) 267 | # ckpt_path = os.path.join(load_directory, "adapter.pt") 268 | # config_path = os.path.join(load_directory, "config.json") 269 | # config = json_load(config_path) 270 | # instance = cls(**config) 271 | # instance.load_state_dict(torch.load(ckpt_path)) 272 | # print(f"Loading adapter model from {ckpt_path}") 273 | # return instance 274 | 275 | 276 | 277 | class ConditionAdapter(nn.Module): 278 | def __init__(self, config): 279 | super(ConditionAdapter, self).__init__() 280 | self.config = config 281 | self.proj = nn.Linear(self.config["condition_dim"], self.config["cross_attention_dim"]) 282 | self.norm = torch.nn.LayerNorm(self.config["cross_attention_dim"]) 283 | print(f"INITIATED: ConditionAdapter: {self.config}") 284 | 285 | def forward(self, x): 286 | x = self.proj(x) 287 | x = self.norm(x) 288 | return x 289 | 290 | @classmethod 291 | def from_pretrained(cls, pretrained_model_name_or_path): 292 | config_path = os.path.join(pretrained_model_name_or_path, "config.json") 293 | ckpt_path = os.path.join(pretrained_model_name_or_path, "condition_adapter.pt") 294 | config = json_load(config_path) 295 | instance = cls(config) 296 | instance.load_state_dict(torch.load(ckpt_path)) 297 | print(f"LOADED: ConditionAdapter from {pretrained_model_name_or_path}") 298 | return instance 299 | 300 | def save_pretrained(self, pretrained_model_name_or_path): 301 | os.makedirs(pretrained_model_name_or_path, exist_ok=True) 302 | config_path = os.path.join(pretrained_model_name_or_path, "config.json") 303 | ckpt_path = os.path.join(pretrained_model_name_or_path, "condition_adapter.pt") 304 | json_dump(self.config, config_path) 305 | torch.save(self.state_dict(), ckpt_path) 306 | print(f"SAVED: ConditionAdapter {self.config['condition_adapter_name']} to {pretrained_model_name_or_path}") 307 | 308 | 309 | # class TextEncoderWrapper(CLIPTextModel): 310 | # def __init__(self, text_encoder, text_encoder_adapter): 311 | # super().__init__(text_encoder.config) 312 | # self.text_encoder = text_encoder 313 | # self.adapter = text_encoder_adapter 314 | 315 | # def forward(self, input_ids, **kwargs): 316 | # outputs = self.text_encoder(input_ids, **kwargs) 317 | # adapted_output = self.adapter(outputs[0]) 318 | # return [adapted_output] # to compatible with last_hidden_state 319 | 320 | -------------------------------------------------------------------------------- /prompt2prompt/ptp_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import numpy as np 16 | import torch 17 | from PIL import Image, ImageDraw, ImageFont 18 | import cv2 19 | from typing import Optional, Union, Tuple, List, Callable, Dict 20 | from tqdm.notebook import tqdm 21 | 22 | 23 | def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0), font_scale: float = 1.0, thickness: int = 3) -> np.ndarray: 24 | h, w, c = image.shape 25 | offset = int(h * .2) 26 | img = np.ones((h + offset, w, c), dtype=np.uint8) * 255 27 | font = cv2.FONT_HERSHEY_SIMPLEX 28 | img[:h] = image 29 | textsize = cv2.getTextSize(text, font, font_scale, thickness)[0] 30 | text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2 31 | cv2.putText(img, text, (text_x, text_y), font, font_scale, text_color, 2) 32 | return img 33 | 34 | 35 | def view_images(images, num_rows=1, offset_ratio=0.02, save_name="null-text+ptp"): 36 | if type(images) is list: 37 | num_empty = len(images) % num_rows 38 | elif images.ndim == 4: 39 | num_empty = images.shape[0] % num_rows 40 | else: 41 | images = [images] 42 | num_empty = 0 43 | 44 | empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255 45 | images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty 46 | num_items = len(images) 47 | 48 | h, w, c = images[0].shape 49 | offset = int(h * offset_ratio) 50 | num_cols = num_items // num_rows 51 | image_ = np.ones((h * num_rows + offset * (num_rows - 1), 52 | w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255 53 | for i in range(num_rows): 54 | for j in range(num_cols): 55 | image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[ 56 | i * num_cols + j] 57 | 58 | pil_img = Image.fromarray(image_).save(f'{save_name}.png') 59 | 60 | def save_image_grid(img, fname, grid_size): 61 | gw, gh = grid_size 62 | _N, C, H, W = img.shape 63 | assert gw * gh == _N 64 | img = img.reshape(gh, gw, C, H, W) 65 | img = img.transpose(0, 3, 1, 4, 2) 66 | img = img.reshape(gh * H, gw * W, C) 67 | 68 | assert C in [1, 3] 69 | if C == 1: 70 | Image.fromarray(img[:, :, 0], 'L').save(fname) 71 | if C == 3: 72 | Image.fromarray(img, 'RGB').save(fname) 73 | 74 | 75 | def diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource=False): 76 | if low_resource: 77 | controller.uncond = True 78 | noise_pred_uncond = model.unet(latents, t, encoder_hidden_states=context[0])["sample"] 79 | controller.uncond = False 80 | noise_prediction_text = model.unet(latents, t, encoder_hidden_states=context[1])["sample"] 81 | else: 82 | latents_input = torch.cat([latents] * 2) 83 | noise_pred = model.unet(latents_input, t, encoder_hidden_states=context)["sample"] 84 | noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2) 85 | noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) 86 | latents = model.scheduler.step(noise_pred, t, latents)["prev_sample"] 87 | latents = controller.step_callback(latents) 88 | return latents 89 | 90 | 91 | def latent2image(vae, latents): 92 | latents = 1 / 0.18215 * latents 93 | image = vae.decode(latents)['sample'] 94 | image = (image / 2 + 0.5).clamp(0, 1) 95 | image = image.cpu().permute(0, 2, 3, 1).numpy() 96 | image = (image * 255).astype(np.uint8) 97 | return image 98 | 99 | 100 | def init_latent(latent, model, height, width, generator, batch_size): 101 | if latent is None: 102 | latent = torch.randn( 103 | (1, model.unet.in_channels, height // 8, width // 8), 104 | generator=generator, 105 | ) 106 | latents = latent.expand(batch_size, model.unet.in_channels, height // 8, width // 8).to(model.device) 107 | return latent, latents 108 | 109 | 110 | @torch.no_grad() 111 | def text2image_ldm( 112 | model, 113 | prompt: List[str], 114 | controller, 115 | num_inference_steps: int = 50, 116 | guidance_scale: Optional[float] = 7., 117 | generator: Optional[torch.Generator] = None, 118 | latent: Optional[torch.FloatTensor] = None, 119 | ): 120 | register_attention_control(model, controller) 121 | height = width = 256 122 | batch_size = len(prompt) 123 | 124 | uncond_input = model.tokenizer([""] * batch_size, padding="max_length", max_length=77, return_tensors="pt") 125 | uncond_embeddings = model.bert(uncond_input.input_ids.to(model.device))[0] 126 | 127 | text_input = model.tokenizer(prompt, padding="max_length", max_length=77, return_tensors="pt") 128 | text_embeddings = model.bert(text_input.input_ids.to(model.device))[0] 129 | latent, latents = init_latent(latent, model, height, width, generator, batch_size) 130 | context = torch.cat([uncond_embeddings, text_embeddings]) 131 | 132 | model.scheduler.set_timesteps(num_inference_steps) 133 | for t in tqdm(model.scheduler.timesteps): 134 | latents = diffusion_step(model, controller, latents, context, t, guidance_scale) 135 | 136 | image = latent2image(model.vqvae, latents) 137 | 138 | return image, latent 139 | 140 | 141 | @torch.no_grad() 142 | def text2image_ldm_stable( 143 | model, 144 | prompt: List[str], 145 | controller, 146 | num_inference_steps: int = 50, 147 | guidance_scale: float = 7.5, 148 | generator: Optional[torch.Generator] = None, 149 | latent: Optional[torch.FloatTensor] = None, 150 | low_resource: bool = False, 151 | ): 152 | register_attention_control(model, controller) 153 | height = width = 512 154 | batch_size = len(prompt) 155 | 156 | text_input = model.tokenizer( 157 | prompt, 158 | padding="max_length", 159 | max_length=model.tokenizer.model_max_length, 160 | truncation=True, 161 | return_tensors="pt", 162 | ) 163 | text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0] 164 | max_length = text_input.input_ids.shape[-1] 165 | uncond_input = model.tokenizer( 166 | [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt" 167 | ) 168 | uncond_embeddings = model.text_encoder(uncond_input.input_ids.to(model.device))[0] 169 | 170 | context = [uncond_embeddings, text_embeddings] 171 | if not low_resource: 172 | context = torch.cat(context) 173 | latent, latents = init_latent(latent, model, height, width, generator, batch_size) 174 | 175 | # set timesteps 176 | extra_set_kwargs = {"offset": 1} 177 | model.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs) 178 | for t in tqdm(model.scheduler.timesteps): 179 | latents = diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource) 180 | 181 | image = latent2image(model.vae, latents) 182 | 183 | return image, latent 184 | 185 | def register_attention_control(model, controller): 186 | def ca_forward(self, place_in_unet): 187 | to_out = self.to_out 188 | if type(to_out) is torch.nn.modules.container.ModuleList: 189 | to_out = self.to_out[0] 190 | else: 191 | to_out = self.to_out 192 | 193 | def forward(hidden_states, encoder_hidden_states=None, attention_mask=None,temb=None,): 194 | is_cross = encoder_hidden_states is not None 195 | 196 | residual = hidden_states 197 | 198 | if self.spatial_norm is not None: 199 | hidden_states = self.spatial_norm(hidden_states, temb) 200 | 201 | input_ndim = hidden_states.ndim 202 | 203 | if input_ndim == 4: 204 | batch_size, channel, height, width = hidden_states.shape 205 | hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) 206 | 207 | batch_size, sequence_length, _ = ( 208 | hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape 209 | ) 210 | attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size) 211 | 212 | if self.group_norm is not None: 213 | hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) 214 | 215 | query = self.to_q(hidden_states) 216 | 217 | if encoder_hidden_states is None: 218 | encoder_hidden_states = hidden_states 219 | elif self.norm_cross: 220 | encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states) 221 | 222 | key = self.to_k(encoder_hidden_states) 223 | value = self.to_v(encoder_hidden_states) 224 | 225 | query = self.head_to_batch_dim(query) 226 | key = self.head_to_batch_dim(key) 227 | value = self.head_to_batch_dim(value) 228 | 229 | attention_probs = self.get_attention_scores(query, key, attention_mask) 230 | attention_probs = controller(attention_probs, is_cross, place_in_unet) 231 | 232 | hidden_states = torch.bmm(attention_probs, value) 233 | hidden_states = self.batch_to_head_dim(hidden_states) 234 | 235 | # linear proj 236 | hidden_states = to_out(hidden_states) 237 | 238 | if input_ndim == 4: 239 | hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) 240 | 241 | if self.residual_connection: 242 | hidden_states = hidden_states + residual 243 | 244 | hidden_states = hidden_states / self.rescale_output_factor 245 | 246 | return hidden_states 247 | 248 | return forward 249 | 250 | class DummyController: 251 | 252 | def __call__(self, *args): 253 | return args[0] 254 | 255 | def __init__(self): 256 | self.num_att_layers = 0 257 | self.ddim_inv = True 258 | 259 | if controller is None: 260 | controller = DummyController() 261 | 262 | def register_recr(net_, count, place_in_unet): 263 | if net_.__class__.__name__ == 'Attention': 264 | net_.forward = ca_forward(net_, place_in_unet) 265 | return count + 1 266 | elif hasattr(net_, 'children'): 267 | for net__ in net_.children(): 268 | count = register_recr(net__, count, place_in_unet) 269 | return count 270 | 271 | cross_att_count = 0 272 | sub_nets = model.unet.named_children() 273 | for net in sub_nets: 274 | if "down" in net[0]: 275 | cross_att_count += register_recr(net[1], 0, "down") 276 | elif "up" in net[0]: 277 | cross_att_count += register_recr(net[1], 0, "up") 278 | elif "mid" in net[0]: 279 | cross_att_count += register_recr(net[1], 0, "mid") 280 | 281 | controller.num_att_layers = cross_att_count 282 | 283 | def get_word_inds(text: str, word_place: int, tokenizer): 284 | split_text = text.split(" ") 285 | if type(word_place) is str: 286 | word_place = [i for i, word in enumerate(split_text) if word_place == word] 287 | elif type(word_place) is int: 288 | word_place = [word_place] 289 | out = [] 290 | if len(word_place) > 0: 291 | words_encode = [tokenizer.decode([item]).strip("#") for item in tokenizer.encode(text)][1:-1] 292 | cur_len, ptr = 0, 0 293 | 294 | for i in range(len(words_encode)): 295 | cur_len += len(words_encode[i]) 296 | if ptr in word_place: 297 | out.append(i + 1) 298 | if cur_len >= len(split_text[ptr]): 299 | ptr += 1 300 | cur_len = 0 301 | return np.array(out) 302 | 303 | 304 | def update_alpha_time_word(alpha, bounds: Union[float, Tuple[float, float]], prompt_ind: int, 305 | word_inds: Optional[torch.Tensor] = None): 306 | if type(bounds) is float: 307 | bounds = 0, bounds 308 | start, end = int(bounds[0] * alpha.shape[0]), int(bounds[1] * alpha.shape[0]) 309 | if word_inds is None: 310 | word_inds = torch.arange(alpha.shape[2]) 311 | alpha[: start, prompt_ind, word_inds] = 0 312 | alpha[start: end, prompt_ind, word_inds] = 1 313 | alpha[end:, prompt_ind, word_inds] = 0 314 | return alpha 315 | 316 | 317 | def get_time_words_attention_alpha(prompts, num_steps, 318 | cross_replace_steps: Union[float, Dict[str, Tuple[float, float]]], 319 | tokenizer, max_num_words=77): 320 | if type(cross_replace_steps) is not dict: 321 | cross_replace_steps = {"default_": cross_replace_steps} 322 | if "default_" not in cross_replace_steps: 323 | cross_replace_steps["default_"] = (0., 1.) 324 | alpha_time_words = torch.zeros(num_steps + 1, len(prompts) - 1, max_num_words) 325 | for i in range(len(prompts) - 1): 326 | alpha_time_words = update_alpha_time_word(alpha_time_words, cross_replace_steps["default_"], 327 | i) 328 | for key, item in cross_replace_steps.items(): 329 | if key != "default_": 330 | inds = [get_word_inds(prompts[i], key, tokenizer) for i in range(1, len(prompts))] 331 | for i, ind in enumerate(inds): 332 | if len(ind) > 0: 333 | alpha_time_words = update_alpha_time_word(alpha_time_words, item, i, ind) 334 | alpha_time_words = alpha_time_words.reshape(num_steps + 1, len(prompts) - 1, 1, 1, max_num_words) 335 | return alpha_time_words 336 | -------------------------------------------------------------------------------- /utils/converter.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | from PIL import Image 5 | 6 | import math 7 | import os 8 | import random 9 | import torch 10 | import json 11 | import torch.utils.data 12 | import numpy as np 13 | import librosa 14 | from librosa.util import normalize 15 | from scipy.io.wavfile import read 16 | from librosa.filters import mel as librosa_mel_fn 17 | 18 | import torch.nn.functional as F 19 | import torch.nn as nn 20 | from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d 21 | from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm 22 | 23 | MAX_WAV_VALUE = 32768.0 24 | 25 | 26 | def load_wav(full_path): 27 | sampling_rate, data = read(full_path) 28 | return data, sampling_rate 29 | 30 | 31 | def dynamic_range_compression(x, C=1, clip_val=1e-5): 32 | return np.log(np.clip(x, a_min=clip_val, a_max=None) * C) 33 | 34 | 35 | def dynamic_range_decompression(x, C=1): 36 | return np.exp(x) / C 37 | 38 | 39 | def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): 40 | return torch.log(torch.clamp(x, min=clip_val) * C) 41 | 42 | 43 | def dynamic_range_decompression_torch(x, C=1): 44 | return torch.exp(x) / C 45 | 46 | 47 | def spectral_normalize_torch(magnitudes): 48 | output = dynamic_range_compression_torch(magnitudes) 49 | return output 50 | 51 | 52 | def spectral_de_normalize_torch(magnitudes): 53 | output = dynamic_range_decompression_torch(magnitudes) 54 | return output 55 | 56 | 57 | mel_basis = {} 58 | hann_window = {} 59 | 60 | 61 | def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False): 62 | if torch.min(y) < -1.: 63 | print('min value is ', torch.min(y)) 64 | if torch.max(y) > 1.: 65 | print('max value is ', torch.max(y)) 66 | 67 | global mel_basis, hann_window 68 | if fmax not in mel_basis: 69 | mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) 70 | mel_basis[str(fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device) 71 | hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device) 72 | 73 | y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect') 74 | y = y.squeeze(1) 75 | 76 | # complex tensor as default, then use view_as_real for future pytorch compatibility 77 | spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[str(y.device)], 78 | center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True) 79 | spec = torch.view_as_real(spec) 80 | spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9)) 81 | 82 | spec = torch.matmul(mel_basis[str(fmax)+'_'+str(y.device)], spec) 83 | spec = spectral_normalize_torch(spec) 84 | 85 | return spec 86 | 87 | 88 | def spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False): 89 | if torch.min(y) < -1.: 90 | print('min value is ', torch.min(y)) 91 | if torch.max(y) > 1.: 92 | print('max value is ', torch.max(y)) 93 | 94 | global hann_window 95 | hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device) 96 | 97 | y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect') 98 | y = y.squeeze(1) 99 | 100 | # complex tensor as default, then use view_as_real for future pytorch compatibility 101 | spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[str(y.device)], 102 | center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True) 103 | spec = torch.view_as_real(spec) 104 | spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9)) 105 | 106 | return spec 107 | 108 | 109 | def normalize_spectrogram( 110 | spectrogram: torch.Tensor, 111 | max_value: float = 200, 112 | min_value: float = 1e-5, 113 | power: float = 1., 114 | inverse: bool = False 115 | ) -> torch.Tensor: 116 | 117 | # Rescale to 0-1 118 | max_value = np.log(max_value) # 5.298317366548036 119 | min_value = np.log(min_value) # -11.512925464970229 120 | 121 | # assert spectrogram.max() <= max_value and spectrogram.min() >= min_value 122 | assert spectrogram.max() <= max_value 123 | assert spectrogram.min() >= min_value 124 | 125 | data = (spectrogram - min_value) / (max_value - min_value) 126 | 127 | # Invert 128 | if inverse: 129 | data = 1 - data 130 | 131 | # Apply the power curve 132 | data = torch.pow(data, power) 133 | 134 | # 1D -> 3D 135 | data = data.repeat(3, 1, 1) 136 | 137 | # Flip Y axis: image origin at the top-left corner, spectrogram origin at the bottom-left corner 138 | data = torch.flip(data, [1]) 139 | 140 | return data 141 | 142 | 143 | 144 | def denormalize_spectrogram( 145 | data: torch.Tensor, 146 | max_value: float = 200, 147 | min_value: float = 1e-5, 148 | power: float = 1, 149 | inverse: bool = False, 150 | ) -> torch.Tensor: 151 | 152 | max_value = np.log(max_value) 153 | min_value = np.log(min_value) 154 | 155 | # Flip Y axis: image origin at the top-left corner, spectrogram origin at the bottom-left corner 156 | data = torch.flip(data, [1]) 157 | 158 | assert len(data.shape) == 3, "Expected 3 dimensions, got {}".format(len(data.shape)) 159 | 160 | if data.shape[0] == 1: 161 | data = data.repeat(3, 1, 1) 162 | 163 | assert data.shape[0] == 3, "Expected 3 channels, got {}".format(data.shape[0]) 164 | data = data[0] 165 | 166 | # Reverse the power curve 167 | data = torch.pow(data, 1 / power) 168 | 169 | # Invert 170 | if inverse: 171 | data = 1 - data 172 | 173 | # Rescale to max value 174 | spectrogram = data * (max_value - min_value) + min_value 175 | 176 | return spectrogram 177 | 178 | 179 | def get_mel_spectrogram_from_audio(audio, device="cuda"): 180 | audio = audio / MAX_WAV_VALUE 181 | audio = librosa.util.normalize(audio) * 0.95 182 | 183 | audio = torch.FloatTensor(audio) 184 | audio = audio.unsqueeze(0) 185 | 186 | if len(audio.shape) == 3: 187 | audio_reduced = audio[:, :, 0] 188 | else: 189 | audio_reduced = audio 190 | # print(f"Processed audio_reduced shape: {audio_reduced.shape}") 191 | target_length = 160000 192 | current_length = audio_reduced.size(1) 193 | if current_length < target_length: 194 | padding = target_length - current_length 195 | audio_reduced = F.pad(audio_reduced, (0, padding), mode='constant', value=0) 196 | elif current_length > target_length: 197 | audio_reduced = audio_reduced[:, :target_length] 198 | # print(f"Processed audio_reduced shape: {audio_reduced.shape}") 199 | 200 | 201 | waveform = audio_reduced.to(device) 202 | spec = mel_spectrogram(waveform, n_fft=2048, num_mels=256, sampling_rate=16000, hop_size=160, win_size=1024, fmin=0, fmax=8000, center=False) 203 | return audio_reduced, spec 204 | 205 | 206 | 207 | LRELU_SLOPE = 0.1 208 | MAX_WAV_VALUE = 32768.0 209 | 210 | 211 | class AttrDict(dict): 212 | def __init__(self, *args, **kwargs): 213 | super(AttrDict, self).__init__(*args, **kwargs) 214 | self.__dict__ = self 215 | 216 | 217 | def get_config(config_path): 218 | config = json.loads(open(config_path).read()) 219 | config = AttrDict(config) 220 | return config 221 | 222 | def init_weights(m, mean=0.0, std=0.01): 223 | classname = m.__class__.__name__ 224 | if classname.find("Conv") != -1: 225 | m.weight.data.normal_(mean, std) 226 | 227 | 228 | def apply_weight_norm(m): 229 | classname = m.__class__.__name__ 230 | if classname.find("Conv") != -1: 231 | weight_norm(m) 232 | 233 | 234 | def get_padding(kernel_size, dilation=1): 235 | return int((kernel_size*dilation - dilation)/2) 236 | 237 | 238 | class ResBlock1(torch.nn.Module): 239 | def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)): 240 | super(ResBlock1, self).__init__() 241 | self.h = h 242 | self.convs1 = nn.ModuleList([ 243 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], 244 | padding=get_padding(kernel_size, dilation[0]))), 245 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], 246 | padding=get_padding(kernel_size, dilation[1]))), 247 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], 248 | padding=get_padding(kernel_size, dilation[2]))) 249 | ]) 250 | self.convs1.apply(init_weights) 251 | 252 | self.convs2 = nn.ModuleList([ 253 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, 254 | padding=get_padding(kernel_size, 1))), 255 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, 256 | padding=get_padding(kernel_size, 1))), 257 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, 258 | padding=get_padding(kernel_size, 1))) 259 | ]) 260 | self.convs2.apply(init_weights) 261 | 262 | def forward(self, x): 263 | for c1, c2 in zip(self.convs1, self.convs2): 264 | xt = F.leaky_relu(x, LRELU_SLOPE) 265 | xt = c1(xt) 266 | xt = F.leaky_relu(xt, LRELU_SLOPE) 267 | xt = c2(xt) 268 | x = xt + x 269 | return x 270 | 271 | def remove_weight_norm(self): 272 | for l in self.convs1: 273 | remove_weight_norm(l) 274 | for l in self.convs2: 275 | remove_weight_norm(l) 276 | 277 | 278 | class ResBlock2(torch.nn.Module): 279 | def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)): 280 | super(ResBlock2, self).__init__() 281 | self.h = h 282 | self.convs = nn.ModuleList([ 283 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], 284 | padding=get_padding(kernel_size, dilation[0]))), 285 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], 286 | padding=get_padding(kernel_size, dilation[1]))) 287 | ]) 288 | self.convs.apply(init_weights) 289 | 290 | def forward(self, x): 291 | for c in self.convs: 292 | xt = F.leaky_relu(x, LRELU_SLOPE) 293 | xt = c(xt) 294 | x = xt + x 295 | return x 296 | 297 | def remove_weight_norm(self): 298 | for l in self.convs: 299 | remove_weight_norm(l) 300 | 301 | 302 | 303 | class Generator(torch.nn.Module): 304 | def __init__(self, h): 305 | super(Generator, self).__init__() 306 | self.h = h 307 | self.num_kernels = len(h.resblock_kernel_sizes) 308 | self.num_upsamples = len(h.upsample_rates) 309 | self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) # change: 80 --> 512 310 | resblock = ResBlock1 if h.resblock == '1' else ResBlock2 311 | 312 | self.ups = nn.ModuleList() 313 | for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): 314 | if (k-u) % 2 == 0: 315 | self.ups.append(weight_norm( 316 | ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)), 317 | k, u, padding=(k-u)//2))) 318 | else: 319 | self.ups.append(weight_norm( 320 | ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)), 321 | k, u, padding=(k-u)//2+1, output_padding=1))) 322 | 323 | 324 | self.resblocks = nn.ModuleList() 325 | for i in range(len(self.ups)): 326 | ch = h.upsample_initial_channel//(2**(i+1)) 327 | for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)): 328 | self.resblocks.append(resblock(h, ch, k, d)) 329 | 330 | self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3)) 331 | self.ups.apply(init_weights) 332 | self.conv_post.apply(init_weights) 333 | 334 | def forward(self, x): 335 | x = self.conv_pre(x) 336 | for i in range(self.num_upsamples): 337 | x = F.leaky_relu(x, LRELU_SLOPE) 338 | x = self.ups[i](x) 339 | xs = None 340 | for j in range(self.num_kernels): 341 | if xs is None: 342 | xs = self.resblocks[i*self.num_kernels+j](x) 343 | else: 344 | xs += self.resblocks[i*self.num_kernels+j](x) 345 | x = xs / self.num_kernels 346 | x = F.leaky_relu(x) 347 | x = self.conv_post(x) 348 | x = torch.tanh(x) 349 | 350 | return x 351 | 352 | def remove_weight_norm(self): 353 | for l in self.ups: 354 | remove_weight_norm(l) 355 | for l in self.resblocks: 356 | l.remove_weight_norm() 357 | remove_weight_norm(self.conv_pre) 358 | remove_weight_norm(self.conv_post) 359 | 360 | @classmethod 361 | def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None): 362 | if subfolder is not None: 363 | pretrained_model_name_or_path = os.path.join(pretrained_model_name_or_path, subfolder) 364 | config_path = os.path.join(pretrained_model_name_or_path, "config.json") 365 | ckpt_path = os.path.join(pretrained_model_name_or_path, "vocoder.pt") 366 | 367 | config = get_config(config_path) 368 | vocoder = cls(config) 369 | 370 | state_dict_g = torch.load(ckpt_path) 371 | vocoder.load_state_dict(state_dict_g["generator"]) 372 | vocoder.eval() 373 | vocoder.remove_weight_norm() 374 | return vocoder 375 | 376 | 377 | @torch.no_grad() 378 | def inference(self, mels, lengths=None): 379 | self.eval() 380 | with torch.no_grad(): 381 | wavs = self(mels).squeeze(1) 382 | 383 | wavs = (wavs.cpu().numpy() * MAX_WAV_VALUE).astype("int16") 384 | 385 | if lengths is not None: 386 | wavs = wavs[:, :lengths] 387 | 388 | return wavs 389 | -------------------------------------------------------------------------------- /eval/configs/audiocaps/edit_add.csv: -------------------------------------------------------------------------------- 1 | Uid,Caption,edit_Caption 2 | jI8-_usilWE_0_10000,Motorcycle engine running,"Motorcycle engine running, and communicating" 3 | dYvKBG-8ha0_30000_40000,Helicopter blades spinning,"Helicopter blades spinning, as fireworks" 4 | ui4PHURhahA_100000_110000,A duck quacking,"A duck quacking, as jet engine" 5 | pOwNRM3BM_w_30000_40000,Water running continuously,"Water running continuously, while animals making noises" 6 | 9RtgQ1ElzfM_30000_40000,Water running from a flushed toilet,"Water running from a flushed toilet, as communicating" 7 | WLH87N00WoM_20000_30000,A man speaks followed by vibrations of a power tool,"A man speaks followed by vibrations of a power tool, and jet engine" 8 | -ftxbusleO8_30000_40000,An engine running,"An engine running, as communicating" 9 | ui4PHURhahA_100000_110000,A duck quacking,"A duck quacking, while train horn" 10 | IHnAMVPBoOA_80000_90000,Church bells ringing,"Church bells ringing, and Music plays" 11 | s1szx3qv84E_10000_20000,A bell is ringing,"A bell is ringing, with raining" 12 | ydO6ZIcghzU_0_10000,A person whistling,"A person whistling, with animals making noises" 13 | USle1qAiOdA_0_10000,A cat meowing twice,"A cat meowing twice, while Wind blowing" 14 | zyI4TfScXFA_30000_40000,Humming of a nearby jet engine,"Humming of a nearby jet engine, and raining" 15 | dpnzm2DPoBI_40000_50000,An infant crying,"An infant crying, with child shouts" 16 | jY2qBSkbZmc_230000_240000,Several loud burps,"Several loud burps, with helicopter flying" 17 | zpMwBUIuxgY_360000_370000,Typing on a computer keyboard,"Typing on a computer keyboard, and fireworks" 18 | F-qkeRTl8gY_0_10000,A vehicle engine starting up then running idle,"A vehicle engine starting up then running idle, with animals making noises" 19 | dpnzm2DPoBI_40000_50000,An infant crying,"An infant crying, and Music plays" 20 | zf9Uk3HjD4k_20000_30000,A telephone ringing,"A telephone ringing, with man talking" 21 | 4hjW385XyHE_30000_40000,A cat meows repeatedly,"A cat meows repeatedly, as child shouts" 22 | cfm3g8IHxcM_20000_30000,People are laughing,"People are laughing, as animals making noises" 23 | ktFiy07FH9s_60000_70000,A toilet is flushing,"A toilet is flushing, and Wind blowing" 24 | Gn2hjm9wx60_50000_60000,A speedboat is traveling across water,"A speedboat is traveling across water, with gun shot" 25 | dpnzm2DPoBI_40000_50000,An infant crying,"An infant crying, while man talking" 26 | 4dyO-l0BFOk_380000_390000,A car is accelerating,"A car is accelerating, while man talking" 27 | lz3At8iTA88_30000_40000,A boat engine running,"A boat engine running, as man talking" 28 | ziUoP8L-zkA_230000_240000,An aircraft engine running,"An aircraft engine running, while sound of an engine" 29 | 0m1sUvpvHu0_30000_40000,Pigeons are cooing,"Pigeons are cooing, and communicating" 30 | zvLcOfhGcGo_0_10000,A person snoring,"A person snoring, and sound of an engine" 31 | yrTqAYH5OM0_30000_40000,Snoring occurs in a rhythmic pattern,"Snoring occurs in a rhythmic pattern, while Wind blowing" 32 | GDMI4uJfDVw_10000_20000,A woman speaking,"A woman speaking, and train horn" 33 | zssgRs7HUmM_30000_40000,A clock ticking,"A clock ticking, and plastic creaks" 34 | Hd_HIXu79jY_30000_40000,A man delivering a speech,"A man delivering a speech, with dog barking" 35 | 0Y1dEz_tVlE_30000_40000,A motor vehicle engine is revving,"A motor vehicle engine is revving, with birds chirping" 36 | w-oSwpsXIuU_20000_30000,Multiple people laughing,"Multiple people laughing, while dog barking" 37 | ypY1zYRcfIc_100000_110000,A toilet flushing,"A toilet flushing, while fireworks" 38 | GDMI4uJfDVw_10000_20000,A woman speaking,"A woman speaking, as dog barking" 39 | zvLcOfhGcGo_0_10000,A person snoring,"A person snoring, as animals making noises" 40 | tZF7dJoT-P0_10000_20000,Continuous burping,"Continuous burping, while sound of an engine" 41 | zvLcOfhGcGo_0_10000,A person snoring,"A person snoring, and child shouts" 42 | zssgRs7HUmM_30000_40000,A clock ticking,"A clock ticking, as animals making noises" 43 | ypY1zYRcfIc_100000_110000,A toilet flushing,"A toilet flushing, and clock chiming" 44 | yWOeusquiJk_20000_30000,A person is snoring,"A person is snoring, and pigeon cooing" 45 | zpMwBUIuxgY_360000_370000,Typing on a computer keyboard,"Typing on a computer keyboard, with train horn" 46 | zO1h2bKYe58_140000_150000,An idle vehicle engine running,"An idle vehicle engine running, with pigeon cooing" 47 | yOpYiaqPwwQ_4000_14000,A baby crying repeatedly,"A baby crying repeatedly, with pigeon cooing" 48 | 91EvnH5ln4Q_30000_40000,A female speaking,"A female speaking, while jet engine" 49 | 8pu1qkMTn-Y_30000_40000,An adult male is speaking,"An adult male is speaking, while communicating" 50 | z8jsXi_J0KM_0_10000,Several goats bleat,"Several goats bleat, with dog barking" 51 | SaUDkGYRZHE_0_10000,A man laughing,"A man laughing, as gun shot" 52 | PzPl-WWTpXs_0_10000,Birds chirping continuously,"Birds chirping continuously, as dog barking" 53 | -ftxbusleO8_30000_40000,An engine running,"An engine running, and helicopter flying" 54 | WE-_ALOv8rI_130000_140000,A man speaks followed by whistling,"A man speaks followed by whistling, and Wind blowing" 55 | -ftxbusleO8_30000_40000,An engine running,"An engine running, and dog barking" 56 | ypY1zYRcfIc_100000_110000,A toilet flushing,"A toilet flushing, while dog barking" 57 | 0V0MxUBtx9g_120000_130000,A motorcycle engine is idling,"A motorcycle engine is idling, with jet engine" 58 | Myma5QTDDO4_110000_120000,A man speaking on a microphone,"A man speaking on a microphone, with raining" 59 | yZpTj9YT-IQ_30000_40000,A man speaking into a microphone,"A man speaking into a microphone, as train horn" 60 | yx6sC-GArA0_3000_13000,A woman speaks followed by a dog barking,"A woman speaks followed by a dog barking, as child shouts" 61 | tNUE8iLhGtg_250000_260000,A toy helicopter flying,"A toy helicopter flying, as gun shot" 62 | sLPAcgyETPM_30000_40000,Female speaking,"Female speaking, as animals making noises" 63 | yc6NJXKEIEk_30000_40000,An engine revving,"An engine revving, and man talking" 64 | z_zkX-8Kf3Q_0_10000,A crowd applauds,"A crowd applauds, while Music plays" 65 | zO1h2bKYe58_140000_150000,An idle vehicle engine running,"An idle vehicle engine running, and pigeon cooing" 66 | FRgQXtYWHhY_30000_40000,A helicopter is in flight,"A helicopter is in flight, as gun shot" 67 | ypY1zYRcfIc_100000_110000,A toilet flushing,"A toilet flushing, while child shouts" 68 | 4XHI4u7VidE_30000_40000,A cat is meowing,"A cat is meowing, and fireworks" 69 | KOO6F0sxAnQ_100000_110000,A man speaking followed by typing on a keyboard,"A man speaking followed by typing on a keyboard, as helicopter flying" 70 | zssgRs7HUmM_30000_40000,A clock ticking,"A clock ticking, with Wind blowing" 71 | xoK7YdYm80I_30000_40000,A helicopter flying,"A helicopter flying, and jet engine" 72 | l2-WXh34EV4_20000_30000,A car engine revving,"A car engine revving, and man talking" 73 | kNGXPjQB9RU_0_10000,A person snores,"A person snores, with train horn" 74 | zpMwBUIuxgY_360000_370000,Typing on a computer keyboard,"Typing on a computer keyboard, while helicopter flying" 75 | dpnzm2DPoBI_40000_50000,An infant crying,"An infant crying, with child shouts" 76 | L77rpdzefmI_30000_40000,Male snoring loudly,"Male snoring loudly, with pigeon cooing" 77 | Amf2hhiHXEo_16000_26000,Engine idling consistently,"Engine idling consistently, and animals making noises" 78 | 84hExVbYonU_30000_40000,A female voice speaking on a microphone,"A female voice speaking on a microphone, as child shouts" 79 | 8eZpe3HQaPk_90000_100000,A helicopter is flying,"A helicopter is flying, as Wind blowing" 80 | pmeh7ocxLP4_0_10000,Continuous snoring,"Continuous snoring, with plastic creaks" 81 | 66a2VZZaY0U_2000_12000,An engine revs several times,"An engine revs several times, with plastic creaks" 82 | zNLca6Dh6-k_120000_130000,Vibrations from a sewing machine,"Vibrations from a sewing machine, and fireworks" 83 | 9285qhAlkZw_30000_40000,A woman gives a speech,"A woman gives a speech, with animals making noises" 84 | zaHl7LdouwY_0_10000,An engine idling continuously,"An engine idling continuously, with communicating" 85 | Bt_nfD_TvcU_30000_40000,A woman delivering a speech,"A woman delivering a speech, and plastic creaks" 86 | ypY1zYRcfIc_100000_110000,A toilet flushing,"A toilet flushing, with animals making noises" 87 | VsfKc_e21WQ_30000_40000,Water running quickly,"Water running quickly, with Wind blowing" 88 | CZFMj95X_sQ_70000_80000,Musical whistling followed by a man speaking,"Musical whistling followed by a man speaking, and gun shot" 89 | 51NfcA7gYzs_20000_30000,Several bells ring,"Several bells ring, while jet engine" 90 | 7pyR815qI4Y_140000_150000,An audience clapping,"An audience clapping, as sound of an engine" 91 | 6C-OgqhhxWk_21000_31000,A motor vehicle engine is idling,"A motor vehicle engine is idling, and Music plays" 92 | yizOr1hPxc8_70000_80000,A helicopter engine running idle,"A helicopter engine running idle, with sound of an engine" 93 | 1mEYMdHfCyc_0_10000,A person types on a keyboard,"A person types on a keyboard, with fireworks" 94 | uEsAup9rS1E_80000_90000,An engine buzzing consistently,"An engine buzzing consistently, as sound of an engine" 95 | 4XHI4u7VidE_30000_40000,A cat is meowing,"A cat is meowing, and man talking" 96 | me05JI_gDac_230000_240000,An engine idling consistently,"An engine idling consistently, and jet engine" 97 | 23cYX9DvH10_3000_13000,A vehicle engine is revving up,"A vehicle engine is revving up, with man talking" 98 | ypY1zYRcfIc_100000_110000,A toilet flushing,"A toilet flushing, while child shouts" 99 | WA9ASASYS8Y_30000_40000,A man speaking,"A man speaking, as dog barking" 100 | zsoRpeRkN34_40000_50000,An aircraft engine sharply running,"An aircraft engine sharply running, and jet engine" 101 | yOpYiaqPwwQ_4000_14000,A baby crying repeatedly,"A baby crying repeatedly, and fireworks" 102 | GDMI4uJfDVw_10000_20000,A woman speaking,"A woman speaking, with clock chiming" 103 | yWOeusquiJk_20000_30000,A person is snoring,"A person is snoring, as plastic creaks" 104 | zssgRs7HUmM_30000_40000,A clock ticking,"A clock ticking, with man talking" 105 | 1IjAxJaJUSk_24000_34000,An engine is vibrating,"An engine is vibrating, while Wind blowing" 106 | 4Ko2TdqMgks_0_10000,Rapid typing on a keyboard,"Rapid typing on a keyboard, and Wind blowing" 107 | yPmy-WCeGKs_0_10000,An engine revving continuously,"An engine revving continuously, as gun shot" 108 | EGve_FMNOwE_30000_40000,A woman is speaking,"A woman is speaking, with clock chiming" 109 | WE-_ALOv8rI_130000_140000,A man speaks followed by whistling,"A man speaks followed by whistling, and dog barking" 110 | 4Jo3aJriewE_30000_40000,A woman delivers a speech,"A woman delivers a speech, as clock chiming" 111 | tNUE8iLhGtg_250000_260000,A toy helicopter flying,"A toy helicopter flying, as train horn" 112 | x9J9pyshnJI_30000_40000,A man speaking continuously,"A man speaking continuously, as man talking" 113 | z2QsWDZMXd0_30000_40000,Ticking of a clock,"Ticking of a clock, with raining" 114 | zvLcOfhGcGo_0_10000,A person snoring,"A person snoring, and sound of an engine" 115 | 6DAUg0daI1I_0_10000,A baby continuously crying,"A baby continuously crying, while woman yells" 116 | tKYiWYC8kiM_290000_300000,Bells ringing a tune,"Bells ringing a tune, as dog barking" 117 | yWOeusquiJk_20000_30000,A person is snoring,"A person is snoring, while plastic creaks" 118 | V7Fcv63Frzc_90000_100000,A man speaks then whistles,"A man speaks then whistles, with dog barking" 119 | ypY1zYRcfIc_100000_110000,A toilet flushing,"A toilet flushing, as fireworks" 120 | 3bNJYlKP5-o_10000_20000,Liquid is being sprayed,"Liquid is being sprayed, while sound of an engine" 121 | zssgRs7HUmM_30000_40000,A clock ticking,"A clock ticking, and pigeon cooing" 122 | -3Tlimy9kJI_240000_250000,A steam engine is hissing,"A steam engine is hissing, while jet engine" 123 | 94rJ6d6WGfk_30000_40000,A young woman speaking,"A young woman speaking, while child shouts" 124 | 6DSQhMvqrTQ_400000_410000,Humming from a large engine,"Humming from a large engine, as animals making noises" 125 | zNLca6Dh6-k_120000_130000,Vibrations from a sewing machine,"Vibrations from a sewing machine, as clock chiming" 126 | zNW0z-SHXTM_70000_80000,Humming of an idling engine,"Humming of an idling engine, as plastic creaks" 127 | yrs-SKTAA9I_80000_90000,Typing on a keyboard,"Typing on a keyboard, with helicopter flying" 128 | lAnEeB55y6g_160000_170000,A car engine idling,"A car engine idling, with train horn" 129 | ylQ3is7jdng_110000_120000,Large church bells ring,"Large church bells ring, with birds chirping" 130 | zO1h2bKYe58_140000_150000,An idle vehicle engine running,"An idle vehicle engine running, with raining" 131 | x9J9pyshnJI_30000_40000,A man speaking continuously,"A man speaking continuously, with fireworks" 132 | zyEGay_RV24_30000_40000,A helicopter flying in the distance,"A helicopter flying in the distance, while helicopter flying" 133 | Bt_nfD_TvcU_30000_40000,A woman delivering a speech,"A woman delivering a speech, as pigeon cooing" 134 | GDMI4uJfDVw_10000_20000,A woman speaking,"A woman speaking, as helicopter flying" 135 | dmMPzURfUw8_480000_490000,"Loud, continuous sizzling","Loud, continuous sizzling, with woman yells" 136 | Eag1nz2oUbk_30000_40000,A woman speaking on a microphone,"A woman speaking on a microphone, with communicating" 137 | yHRLqGdgc6I_40000_50000,A person snoring repetitively,"A person snoring repetitively, as sound of an engine" 138 | 5er9MoTAogc_10000_20000,A baby crying loudly,"A baby crying loudly, as communicating" 139 | GEBMPrzW0NI_30000_40000,A man is speaking,"A man is speaking, with woman yells" 140 | zyEGay_RV24_30000_40000,A helicopter flying in the distance,"A helicopter flying in the distance, with fireworks" 141 | yc6NJXKEIEk_30000_40000,An engine revving,"An engine revving, and child shouts" 142 | 1760Zs-Rezw_20000_30000,A goat bleating,"A goat bleating, with plastic creaks" 143 | 91EvnH5ln4Q_30000_40000,A female speaking,"A female speaking, while jet engine" 144 | z3gvS28om7o_180000_190000,A woman giving a speech,"A woman giving a speech, and jet engine" 145 | 7yYAqtNqM78_0_10000,A person burps loudly,"A person burps loudly, with jet engine" 146 | jY2qBSkbZmc_230000_240000,Several loud burps,"Several loud burps, with gun shot" 147 | zsvaHUeivRo_170000_180000,A power tool drilling,"A power tool drilling, while raining" 148 | yrs-SKTAA9I_80000_90000,Typing on a keyboard,"Typing on a keyboard, while Wind blowing" 149 | GDMI4uJfDVw_10000_20000,A woman speaking,"A woman speaking, while Wind blowing" 150 | U5pCL_ogYuc_30000_40000,An idle motorboat engine running,"An idle motorboat engine running, and birds chirping" 151 | WA9ASASYS8Y_30000_40000,A man speaking,"A man speaking, as helicopter flying" 152 | -M2w9Xc4MUw_1000_11000,Tires squealing followed by an engine revving,"Tires squealing followed by an engine revving, and fireworks" 153 | Bt_nfD_TvcU_30000_40000,A woman delivering a speech,"A woman delivering a speech, and Wind blowing" 154 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # AudioEditor: A Training-Free Diffusion-Based Audio Editing Framework 2 | 3 | [Demo Page](https://kiri0824.github.io/AudioEditor-demo-page/) | [Arxiv Paper](https://arxiv.org/pdf/2409.12466) 4 | 5 | Diffusion-based text-to-audio (TTA) generation has made substantial progress, leveraging Latent Diffusion Model (LDM) to produce high-quality, diverse and instruction-relevant audios. However, beyond generation, the task of audio editing remains equally important but has received comparatively little attention. Audio editing tasks face two primary challenges: executing precise edits and preserving the unedited sections. While workflows based on LDMs have effectively addressed these challenges in the field of image processing, similar approaches have been scarcely applied to audio editing. In this paper, we introduce AudioEditor, a training-free audio editing framework built on the pretrained diffusion-based TTA model. AudioEditor incorporates Null-text Inversion and EOT-Suppression methods, enabling the model to preserve original audio features while executing accurate edits. Comprehensive objective and subjective experiments validate the effectiveness of AudioEditor in delivering high-quality audio edits. 6 | 7 |
8 | Workflow 9 |
10 | 11 | ## 🚀 Features 12 | 13 | - **Pre-trained Model: Auffusion** 14 | We use the pre-trained model Auffusion for audio editing tasks. 15 | [Auffusion Repository](https://github.com/happylittlecat2333/Auffusion/tree/main) | [Model Download Link](https://huggingface.co/auffusion/auffusion) 16 | - **Null-text Inversion**: Ensures preservation of unedited audio portions during the editing process. 17 | - **EOT-suppression**: Enhance the model's ability to preserve original audio 18 | features and improve editing capabilities. 19 | - **Support multiple audio editing operations**: Add, Delete and Replace. 20 | - **Easy integration with other TTA models**: Plug and play with existing TTA diffuser-based models. 21 | 22 | ## 📀 Installation 23 | 24 | 1. Clone the repository: 25 | 26 | ```bash 27 | git clone https://github.com/YuuhangJia/AudioEditor.git 28 | cd AudioEditor 29 | ``` 30 | 31 | 2. Install the required dependencies: 32 | 33 | ```bash 34 | pip install -r requirements.txt 35 | ``` 36 | 37 | 3. Set up additional environment variables or configurations (if any): 38 | 39 | ```bash 40 | export YOUR_ENV_VAR=your_value 41 | ``` 42 | 43 | ## ⭐ Usage 44 | We will release the code for the main methods proposed in our paper after its acceptance immediately! 45 | 46 | ### 1️⃣ Delete 47 | 48 | 49 | Delete/Add Workflow 50 | 51 | To run the deletion on an audio with a simple example, use the following command: 52 | ```bash 53 | python main.py --prompt "After a gunshot, there was a burst of dog barking" \ 54 | --audio_path "audio_examples/input_audios/After a gunshot, there was a burst of dog barking.wav" \ 55 | --token_indices "[[10,11]]" \ 56 | --alpha "[1.,]" --cross_retain_steps "[.2,]" 57 | ``` 58 | 59 | ### 2️⃣ Replace 60 | 61 | Delete/Add Workflow 62 | 63 | To run the Replacement on an audio with a simple example, use the following command: 64 | ```bash 65 | python main.py --prompt "After a thunder, there was a burst of dog barking" \ 66 | --audio_path "audio_examples/input_audios/After a gunshot, there was a burst of dog barking.wav" \ 67 | --token_indices "[[3]]" \ 68 | --alpha "[-0.001,]" --cross_retain_steps "[.2,]" 69 | ``` 70 | 71 | ### 3️⃣ Add 72 | Delete/Add Workflow 73 | 74 | ```bash 75 | python main.py --prompt "A woman is giving a speech amid applause" \ 76 | --audio_path "audio_examples/input_audios/A woman is giving a speech.wav" \ 77 | --token_indices "[[7,8]]" \ 78 | --alpha "[-0.001,]" --cross_retain_steps "[.2,]" 79 | ``` 80 | ## 📐 Quantitative comparison 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | 139 | 140 | 141 | 142 | 143 | 144 | 145 | 146 | 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | 160 | 161 | 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | 211 | 212 | 213 | 214 | 215 | 216 | 217 | 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 | 234 | 235 | 236 | 237 | 238 | 239 | 240 | 241 | 242 | 243 | 244 | 245 | 246 | 247 | 248 | 249 | 250 | 251 | 252 | 253 | 254 | 255 | 256 | 257 | 258 | 259 | 260 | 261 | 262 | 263 | 264 | 265 | 266 | 267 | 268 | 269 | 270 | 271 | 272 | 273 | 274 | 275 | 276 | 277 | 278 | 279 | 280 | 281 | 282 | 283 | 284 | 285 | 286 | 287 | 288 | 289 | 290 |
Objective Evaluation Results
Edit_ModelsEdit_TypeOverall QualitySimilarity with (Regenerated_wavs)Similarity with (Original_wavs)
Clap↑IS ↑FD ↓FAD ↓KL ↓FD ↓FAD ↓KL ↓
Original_wavsadd51.4%5.6444.715.281.78---
delete51.5%4.2651.826.161.85---
replace41.6%4.4169.927.884.56---
Average48.2%4.7755.486.452.73---
Regenerated_wavsadd59.7%5.96---44.715.281.36
delete59.1%4.47---51.826.162.39
replace58.9%5.13---69.927.884.09
Average59.2%5.19---55.486.452.61
SDEdit(baseline)add58.4%6.3627.892.740.7936.743.081.08
delete53.3%5.3155.126.651.7840.436.950.88
replace58.6%4.9929.763.240.8055.217.003.40
Average56.8%5.5537.594.21*1.12*44.13*5.68*1.79
AudioEditor(ours)add59.4%6.1627.832.410.8540.003.521.27
delete54.1%4.7552.565.021.5437.164.911.05
replace58.1%5.1428.803.340.7959.467.523.73
Average57.6%*5.19*37.63*3.271.0743.484.951.93*
291 | 292 | * indicates a suboptimal value, which may represent more desirable than optimal one in certain metrics. 293 | 294 | 295 | 296 | ## 🤝🏻 Contact 297 | Should you have any questions, please contact 2120240729@mail.nankai.edu.cn 298 | 299 | ## 📚 Citation 300 | Coming soon. 301 | 302 | ## 🐍 License 303 | The code in this repository is licensed under the MIT License for academic and other non-commercial uses. 304 | 305 | [//]: # (For commercial use of the code and models, separate commercial licensing is available. Please contact authors.) 306 | 307 | ![visitors](https://visitor-badge.laobi.icu/badge?page_id=sen-mao/SuppressEOT) 308 | 309 | ## 🙏 Acknowledgment: 310 | This code is based on the [P2P, Null-text](https://github.com/google/prompt-to-prompt) , [SuppressEOT](https://github.com/sen-mao/SuppressEOT) and [Auffusion](https://github.com/happylittlecat2333/Auffusion) repositories. 311 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Attribution-NonCommercial-ShareAlike 4.0 International 2 | 3 | ======================================================================= 4 | 5 | Creative Commons Corporation ("Creative Commons") is not a law firm and 6 | does not provide legal services or legal advice. Distribution of 7 | Creative Commons public licenses does not create a lawyer-client or 8 | other relationship. Creative Commons makes its licenses and related 9 | information available on an "as-is" basis. Creative Commons gives no 10 | warranties regarding its licenses, any material licensed under their 11 | terms and conditions, or any related information. Creative Commons 12 | disclaims all liability for damages resulting from their use to the 13 | fullest extent possible. 14 | 15 | Using Creative Commons Public Licenses 16 | 17 | Creative Commons public licenses provide a standard set of terms and 18 | conditions that creators and other rights holders may use to share 19 | original works of authorship and other material subject to copyright 20 | and certain other rights specified in the public license below. The 21 | following considerations are for informational purposes only, are not 22 | exhaustive, and do not form part of our licenses. 23 | 24 | Considerations for licensors: Our public licenses are 25 | intended for use by those authorized to give the public 26 | permission to use material in ways otherwise restricted by 27 | copyright and certain other rights. Our licenses are 28 | irrevocable. Licensors should read and understand the terms 29 | and conditions of the license they choose before applying it. 30 | Licensors should also secure all rights necessary before 31 | applying our licenses so that the public can reuse the 32 | material as expected. Licensors should clearly mark any 33 | material not subject to the license. This includes other CC- 34 | licensed material, or material used under an exception or 35 | limitation to copyright. More considerations for licensors: 36 | wiki.creativecommons.org/Considerations_for_licensors 37 | 38 | Considerations for the public: By using one of our public 39 | licenses, a licensor grants the public permission to use the 40 | licensed material under specified terms and conditions. If 41 | the licensor's permission is not necessary for any reason--for 42 | example, because of any applicable exception or limitation to 43 | copyright--then that use is not regulated by the license. Our 44 | licenses grant only permissions under copyright and certain 45 | other rights that a licensor has authority to grant. Use of 46 | the licensed material may still be restricted for other 47 | reasons, including because others have copyright or other 48 | rights in the material. A licensor may make special requests, 49 | such as asking that all changes be marked or described. 50 | Although not required by our licenses, you are encouraged to 51 | respect those requests where reasonable. More considerations 52 | for the public: 53 | wiki.creativecommons.org/Considerations_for_licensees 54 | 55 | ======================================================================= 56 | 57 | Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 58 | Public License 59 | 60 | By exercising the Licensed Rights (defined below), You accept and agree 61 | to be bound by the terms and conditions of this Creative Commons 62 | Attribution-NonCommercial-ShareAlike 4.0 International Public License 63 | ("Public License"). To the extent this Public License may be 64 | interpreted as a contract, You are granted the Licensed Rights in 65 | consideration of Your acceptance of these terms and conditions, and the 66 | Licensor grants You such rights in consideration of benefits the 67 | Licensor receives from making the Licensed Material available under 68 | these terms and conditions. 69 | 70 | 71 | Section 1 -- Definitions. 72 | 73 | a. Adapted Material means material subject to Copyright and Similar 74 | Rights that is derived from or based upon the Licensed Material 75 | and in which the Licensed Material is translated, altered, 76 | arranged, transformed, or otherwise modified in a manner requiring 77 | permission under the Copyright and Similar Rights held by the 78 | Licensor. For purposes of this Public License, where the Licensed 79 | Material is a musical work, performance, or sound recording, 80 | Adapted Material is always produced where the Licensed Material is 81 | synched in timed relation with a moving image. 82 | 83 | b. Adapter's License means the license You apply to Your Copyright 84 | and Similar Rights in Your contributions to Adapted Material in 85 | accordance with the terms and conditions of this Public License. 86 | 87 | c. BY-NC-SA Compatible License means a license listed at 88 | creativecommons.org/compatiblelicenses, approved by Creative 89 | Commons as essentially the equivalent of this Public License. 90 | 91 | d. Copyright and Similar Rights means copyright and/or similar rights 92 | closely related to copyright including, without limitation, 93 | performance, broadcast, sound recording, and Sui Generis Database 94 | Rights, without regard to how the rights are labeled or 95 | categorized. For purposes of this Public License, the rights 96 | specified in Section 2(b)(1)-(2) are not Copyright and Similar 97 | Rights. 98 | 99 | e. Effective Technological Measures means those measures that, in the 100 | absence of proper authority, may not be circumvented under laws 101 | fulfilling obligations under Article 11 of the WIPO Copyright 102 | Treaty adopted on December 20, 1996, and/or similar international 103 | agreements. 104 | 105 | f. Exceptions and Limitations means fair use, fair dealing, and/or 106 | any other exception or limitation to Copyright and Similar Rights 107 | that applies to Your use of the Licensed Material. 108 | 109 | g. License Elements means the license attributes listed in the name 110 | of a Creative Commons Public License. The License Elements of this 111 | Public License are Attribution, NonCommercial, and ShareAlike. 112 | 113 | h. Licensed Material means the artistic or literary work, database, 114 | or other material to which the Licensor applied this Public 115 | License. 116 | 117 | i. Licensed Rights means the rights granted to You subject to the 118 | terms and conditions of this Public License, which are limited to 119 | all Copyright and Similar Rights that apply to Your use of the 120 | Licensed Material and that the Licensor has authority to license. 121 | 122 | j. Licensor means the individual(s) or entity(ies) granting rights 123 | under this Public License. 124 | 125 | k. NonCommercial means not primarily intended for or directed towards 126 | commercial advantage or monetary compensation. For purposes of 127 | this Public License, the exchange of the Licensed Material for 128 | other material subject to Copyright and Similar Rights by digital 129 | file-sharing or similar means is NonCommercial provided there is 130 | no payment of monetary compensation in connection with the 131 | exchange. 132 | 133 | l. Share means to provide material to the public by any means or 134 | process that requires permission under the Licensed Rights, such 135 | as reproduction, public display, public performance, distribution, 136 | dissemination, communication, or importation, and to make material 137 | available to the public including in ways that members of the 138 | public may access the material from a place and at a time 139 | individually chosen by them. 140 | 141 | m. Sui Generis Database Rights means rights other than copyright 142 | resulting from Directive 96/9/EC of the European Parliament and of 143 | the Council of 11 March 1996 on the legal protection of databases, 144 | as amended and/or succeeded, as well as other essentially 145 | equivalent rights anywhere in the world. 146 | 147 | n. You means the individual or entity exercising the Licensed Rights 148 | under this Public License. Your has a corresponding meaning. 149 | 150 | 151 | Section 2 -- Scope. 152 | 153 | a. License grant. 154 | 155 | 1. Subject to the terms and conditions of this Public License, 156 | the Licensor hereby grants You a worldwide, royalty-free, 157 | non-sublicensable, non-exclusive, irrevocable license to 158 | exercise the Licensed Rights in the Licensed Material to: 159 | 160 | a. reproduce and Share the Licensed Material, in whole or 161 | in part, for NonCommercial purposes only; and 162 | 163 | b. produce, reproduce, and Share Adapted Material for 164 | NonCommercial purposes only. 165 | 166 | 2. Exceptions and Limitations. For the avoidance of doubt, where 167 | Exceptions and Limitations apply to Your use, this Public 168 | License does not apply, and You do not need to comply with 169 | its terms and conditions. 170 | 171 | 3. Term. The term of this Public License is specified in Section 172 | 6(a). 173 | 174 | 4. Media and formats; technical modifications allowed. The 175 | Licensor authorizes You to exercise the Licensed Rights in 176 | all media and formats whether now known or hereafter created, 177 | and to make technical modifications necessary to do so. The 178 | Licensor waives and/or agrees not to assert any right or 179 | authority to forbid You from making technical modifications 180 | necessary to exercise the Licensed Rights, including 181 | technical modifications necessary to circumvent Effective 182 | Technological Measures. For purposes of this Public License, 183 | simply making modifications authorized by this Section 2(a) 184 | (4) never produces Adapted Material. 185 | 186 | 5. Downstream recipients. 187 | 188 | a. Offer from the Licensor -- Licensed Material. Every 189 | recipient of the Licensed Material automatically 190 | receives an offer from the Licensor to exercise the 191 | Licensed Rights under the terms and conditions of this 192 | Public License. 193 | 194 | b. Additional offer from the Licensor -- Adapted Material. 195 | Every recipient of Adapted Material from You 196 | automatically receives an offer from the Licensor to 197 | exercise the Licensed Rights in the Adapted Material 198 | under the conditions of the Adapter's License You apply. 199 | 200 | c. No downstream restrictions. You may not offer or impose 201 | any additional or different terms or conditions on, or 202 | apply any Effective Technological Measures to, the 203 | Licensed Material if doing so restricts exercise of the 204 | Licensed Rights by any recipient of the Licensed 205 | Material. 206 | 207 | 6. No endorsement. Nothing in this Public License constitutes or 208 | may be construed as permission to assert or imply that You 209 | are, or that Your use of the Licensed Material is, connected 210 | with, or sponsored, endorsed, or granted official status by, 211 | the Licensor or others designated to receive attribution as 212 | provided in Section 3(a)(1)(A)(i). 213 | 214 | b. Other rights. 215 | 216 | 1. Moral rights, such as the right of integrity, are not 217 | licensed under this Public License, nor are publicity, 218 | privacy, and/or other similar personality rights; however, to 219 | the extent possible, the Licensor waives and/or agrees not to 220 | assert any such rights held by the Licensor to the limited 221 | extent necessary to allow You to exercise the Licensed 222 | Rights, but not otherwise. 223 | 224 | 2. Patent and trademark rights are not licensed under this 225 | Public License. 226 | 227 | 3. To the extent possible, the Licensor waives any right to 228 | collect royalties from You for the exercise of the Licensed 229 | Rights, whether directly or through a collecting society 230 | under any voluntary or waivable statutory or compulsory 231 | licensing scheme. In all other cases the Licensor expressly 232 | reserves any right to collect such royalties, including when 233 | the Licensed Material is used other than for NonCommercial 234 | purposes. 235 | 236 | 237 | Section 3 -- License Conditions. 238 | 239 | Your exercise of the Licensed Rights is expressly made subject to the 240 | following conditions. 241 | 242 | a. Attribution. 243 | 244 | 1. If You Share the Licensed Material (including in modified 245 | form), You must: 246 | 247 | a. retain the following if it is supplied by the Licensor 248 | with the Licensed Material: 249 | 250 | i. identification of the creator(s) of the Licensed 251 | Material and any others designated to receive 252 | attribution, in any reasonable manner requested by 253 | the Licensor (including by pseudonym if 254 | designated); 255 | 256 | ii. a copyright notice; 257 | 258 | iii. a notice that refers to this Public License; 259 | 260 | iv. a notice that refers to the disclaimer of 261 | warranties; 262 | 263 | v. a URI or hyperlink to the Licensed Material to the 264 | extent reasonably practicable; 265 | 266 | b. indicate if You modified the Licensed Material and 267 | retain an indication of any previous modifications; and 268 | 269 | c. indicate the Licensed Material is licensed under this 270 | Public License, and include the text of, or the URI or 271 | hyperlink to, this Public License. 272 | 273 | 2. You may satisfy the conditions in Section 3(a)(1) in any 274 | reasonable manner based on the medium, means, and context in 275 | which You Share the Licensed Material. For example, it may be 276 | reasonable to satisfy the conditions by providing a URI or 277 | hyperlink to a resource that includes the required 278 | information. 279 | 3. If requested by the Licensor, You must remove any of the 280 | information required by Section 3(a)(1)(A) to the extent 281 | reasonably practicable. 282 | 283 | b. ShareAlike. 284 | 285 | In addition to the conditions in Section 3(a), if You Share 286 | Adapted Material You produce, the following conditions also apply. 287 | 288 | 1. The Adapter's License You apply must be a Creative Commons 289 | license with the same License Elements, this version or 290 | later, or a BY-NC-SA Compatible License. 291 | 292 | 2. You must include the text of, or the URI or hyperlink to, the 293 | Adapter's License You apply. You may satisfy this condition 294 | in any reasonable manner based on the medium, means, and 295 | context in which You Share Adapted Material. 296 | 297 | 3. You may not offer or impose any additional or different terms 298 | or conditions on, or apply any Effective Technological 299 | Measures to, Adapted Material that restrict exercise of the 300 | rights granted under the Adapter's License You apply. 301 | 302 | 303 | Section 4 -- Sui Generis Database Rights. 304 | 305 | Where the Licensed Rights include Sui Generis Database Rights that 306 | apply to Your use of the Licensed Material: 307 | 308 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right 309 | to extract, reuse, reproduce, and Share all or a substantial 310 | portion of the contents of the database for NonCommercial purposes 311 | only; 312 | 313 | b. if You include all or a substantial portion of the database 314 | contents in a database in which You have Sui Generis Database 315 | Rights, then the database in which You have Sui Generis Database 316 | Rights (but not its individual contents) is Adapted Material, 317 | including for purposes of Section 3(b); and 318 | 319 | c. You must comply with the conditions in Section 3(a) if You Share 320 | all or a substantial portion of the contents of the database. 321 | 322 | For the avoidance of doubt, this Section 4 supplements and does not 323 | replace Your obligations under this Public License where the Licensed 324 | Rights include other Copyright and Similar Rights. 325 | 326 | 327 | Section 5 -- Disclaimer of Warranties and Limitation of Liability. 328 | 329 | a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE 330 | EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS 331 | AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF 332 | ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, 333 | IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, 334 | WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR 335 | PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, 336 | ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT 337 | KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT 338 | ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. 339 | 340 | b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE 341 | TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, 342 | NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, 343 | INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, 344 | COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR 345 | USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN 346 | ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR 347 | DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR 348 | IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. 349 | 350 | c. The disclaimer of warranties and limitation of liability provided 351 | above shall be interpreted in a manner that, to the extent 352 | possible, most closely approximates an absolute disclaimer and 353 | waiver of all liability. 354 | 355 | 356 | Section 6 -- Term and Termination. 357 | 358 | a. This Public License applies for the term of the Copyright and 359 | Similar Rights licensed here. However, if You fail to comply with 360 | this Public License, then Your rights under this Public License 361 | terminate automatically. 362 | 363 | b. Where Your right to use the Licensed Material has terminated under 364 | Section 6(a), it reinstates: 365 | 366 | 1. automatically as of the date the violation is cured, provided 367 | it is cured within 30 days of Your discovery of the 368 | violation; or 369 | 370 | 2. upon express reinstatement by the Licensor. 371 | 372 | For the avoidance of doubt, this Section 6(b) does not affect any 373 | right the Licensor may have to seek remedies for Your violations 374 | of this Public License. 375 | 376 | c. For the avoidance of doubt, the Licensor may also offer the 377 | Licensed Material under separate terms or conditions or stop 378 | distributing the Licensed Material at any time; however, doing so 379 | will not terminate this Public License. 380 | 381 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public 382 | License. 383 | 384 | 385 | Section 7 -- Other Terms and Conditions. 386 | 387 | a. The Licensor shall not be bound by any additional or different 388 | terms or conditions communicated by You unless expressly agreed. 389 | 390 | b. Any arrangements, understandings, or agreements regarding the 391 | Licensed Material not stated herein are separate from and 392 | independent of the terms and conditions of this Public License. 393 | 394 | 395 | Section 8 -- Interpretation. 396 | 397 | a. For the avoidance of doubt, this Public License does not, and 398 | shall not be interpreted to, reduce, limit, restrict, or impose 399 | conditions on any use of the Licensed Material that could lawfully 400 | be made without permission under this Public License. 401 | 402 | b. To the extent possible, if any provision of this Public License is 403 | deemed unenforceable, it shall be automatically reformed to the 404 | minimum extent necessary to make it enforceable. If the provision 405 | cannot be reformed, it shall be severed from this Public License 406 | without affecting the enforceability of the remaining terms and 407 | conditions. 408 | 409 | c. No term or condition of this Public License will be waived and no 410 | failure to comply consented to unless expressly agreed to by the 411 | Licensor. 412 | 413 | d. Nothing in this Public License constitutes or may be interpreted 414 | as a limitation upon, or waiver of, any privileges and immunities 415 | that apply to the Licensor or You, including from the legal 416 | processes of any jurisdiction or authority. 417 | 418 | ======================================================================= 419 | 420 | Creative Commons is not a party to its public 421 | licenses. Notwithstanding, Creative Commons may elect to apply one of 422 | its public licenses to material it publishes and in those instances 423 | will be considered the 鈥淟icensor.鈥� The text of the Creative Commons 424 | public licenses is dedicated to the public domain under the CC0 Public 425 | Domain Dedication. Except for the limited purpose of indicating that 426 | material is shared under a Creative Commons public license or as 427 | otherwise permitted by the Creative Commons policies published at 428 | creativecommons.org/policies, Creative Commons does not authorize the 429 | use of the trademark "Creative Commons" or any other trademark or logo 430 | of Creative Commons without its prior written consent including, 431 | without limitation, in connection with any unauthorized modifications 432 | to any of its public licenses or any other arrangements, 433 | understandings, or agreements concerning use of licensed material. For 434 | the avoidance of doubt, this paragraph does not form part of the 435 | public licenses. 436 | 437 | Creative Commons may be contacted at creativecommons.org. 438 | -------------------------------------------------------------------------------- /auffusion/auffusion_pipeline.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 The HuggingFace Team. All rights reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import inspect 16 | import warnings 17 | from typing import Any, Callable, Dict, List, Optional, Union 18 | from dataclasses import dataclass 19 | 20 | import torch 21 | from packaging import version 22 | from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer 23 | 24 | from diffusers.configuration_utils import FrozenDict 25 | from diffusers.image_processor import VaeImageProcessor 26 | from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin 27 | from diffusers.models import AutoencoderKL, UNet2DConditionModel 28 | from diffusers.schedulers import KarrasDiffusionSchedulers 29 | from diffusers.utils import ( 30 | deprecate, 31 | is_accelerate_available, 32 | is_accelerate_version, 33 | logging, 34 | randn_tensor, 35 | replace_example_docstring, 36 | ) 37 | from diffusers.pipelines.pipeline_utils import DiffusionPipeline 38 | from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput 39 | from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker 40 | from huggingface_hub import snapshot_download 41 | from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel, PNDMScheduler 42 | from transformers import PretrainedConfig, AutoTokenizer 43 | import torch.nn as nn 44 | import os, json, PIL 45 | import numpy as np 46 | import torch.nn.functional as F 47 | from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d 48 | from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm 49 | from diffusers.utils.outputs import BaseOutput 50 | import matplotlib.pyplot as plt 51 | 52 | from prompt2prompt.attn_control import AttentionStore 53 | 54 | 55 | logger = logging.get_logger(__name__) # pylint: disable=invalid-name 56 | 57 | 58 | 59 | def json_dump(data_json, json_save_path): 60 | with open(json_save_path, 'w') as f: 61 | json.dump(data_json, f, indent=4) 62 | f.close() 63 | 64 | 65 | def json_load(json_path): 66 | with open(json_path, 'r') as f: 67 | data = json.load(f) 68 | f.close() 69 | return data 70 | 71 | 72 | def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str): 73 | text_encoder_config = PretrainedConfig.from_pretrained( 74 | pretrained_model_name_or_path 75 | ) 76 | model_class = text_encoder_config.architectures[0] 77 | 78 | if model_class == "CLIPTextModel": 79 | from transformers import CLIPTextModel 80 | return CLIPTextModel 81 | if "t5" in model_class.lower(): 82 | from transformers import T5EncoderModel 83 | return T5EncoderModel 84 | if "clap" in model_class.lower(): 85 | from transformers import ClapTextModelWithProjection 86 | return ClapTextModelWithProjection 87 | else: 88 | raise ValueError(f"{model_class} is not supported.") 89 | 90 | 91 | class ConditionAdapter(nn.Module): 92 | def __init__(self, config): 93 | super(ConditionAdapter, self).__init__() 94 | self.config = config 95 | self.proj = nn.Linear(self.config["condition_dim"], self.config["cross_attention_dim"]) 96 | self.norm = torch.nn.LayerNorm(self.config["cross_attention_dim"]) 97 | print(f"INITIATED: ConditionAdapter: {self.config}") 98 | 99 | def forward(self, x): 100 | x = self.proj(x) 101 | x = self.norm(x) 102 | return x 103 | 104 | @classmethod 105 | def from_pretrained(cls, pretrained_model_name_or_path): 106 | config_path = os.path.join(pretrained_model_name_or_path, "config.json") 107 | ckpt_path = os.path.join(pretrained_model_name_or_path, "condition_adapter.pt") 108 | config = json.loads(open(config_path).read()) 109 | instance = cls(config) 110 | instance.load_state_dict(torch.load(ckpt_path)) 111 | print(f"LOADED: ConditionAdapter from {pretrained_model_name_or_path}") 112 | return instance 113 | 114 | def save_pretrained(self, pretrained_model_name_or_path): 115 | os.makedirs(pretrained_model_name_or_path, exist_ok=True) 116 | config_path = os.path.join(pretrained_model_name_or_path, "config.json") 117 | ckpt_path = os.path.join(pretrained_model_name_or_path, "condition_adapter.pt") 118 | json_dump(self.config, config_path) 119 | torch.save(self.state_dict(), ckpt_path) 120 | print(f"SAVED: ConditionAdapter {self.config['model_name']} to {pretrained_model_name_or_path}") 121 | 122 | 123 | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): 124 | """ 125 | Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and 126 | Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 127 | """ 128 | std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) 129 | std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) 130 | # rescale the results from guidance (fixes overexposure) 131 | noise_pred_rescaled = noise_cfg * (std_text / std_cfg) 132 | # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images 133 | noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg 134 | return noise_cfg 135 | 136 | 137 | 138 | LRELU_SLOPE = 0.1 139 | MAX_WAV_VALUE = 32768.0 140 | 141 | 142 | class AttrDict(dict): 143 | def __init__(self, *args, **kwargs): 144 | super(AttrDict, self).__init__(*args, **kwargs) 145 | self.__dict__ = self 146 | 147 | def get_config(config_path): 148 | config = json.loads(open(config_path).read()) 149 | config = AttrDict(config) 150 | return config 151 | 152 | def init_weights(m, mean=0.0, std=0.01): 153 | classname = m.__class__.__name__ 154 | if classname.find("Conv") != -1: 155 | m.weight.data.normal_(mean, std) 156 | 157 | 158 | def apply_weight_norm(m): 159 | classname = m.__class__.__name__ 160 | if classname.find("Conv") != -1: 161 | weight_norm(m) 162 | 163 | 164 | def get_padding(kernel_size, dilation=1): 165 | return int((kernel_size*dilation - dilation)/2) 166 | 167 | 168 | class ResBlock1(torch.nn.Module): 169 | def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)): 170 | super(ResBlock1, self).__init__() 171 | self.h = h 172 | self.convs1 = nn.ModuleList([ 173 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], 174 | padding=get_padding(kernel_size, dilation[0]))), 175 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], 176 | padding=get_padding(kernel_size, dilation[1]))), 177 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], 178 | padding=get_padding(kernel_size, dilation[2]))) 179 | ]) 180 | self.convs1.apply(init_weights) 181 | 182 | self.convs2 = nn.ModuleList([ 183 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, 184 | padding=get_padding(kernel_size, 1))), 185 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, 186 | padding=get_padding(kernel_size, 1))), 187 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, 188 | padding=get_padding(kernel_size, 1))) 189 | ]) 190 | self.convs2.apply(init_weights) 191 | 192 | def forward(self, x): 193 | for c1, c2 in zip(self.convs1, self.convs2): 194 | xt = F.leaky_relu(x, LRELU_SLOPE) 195 | xt = c1(xt) 196 | xt = F.leaky_relu(xt, LRELU_SLOPE) 197 | xt = c2(xt) 198 | x = xt + x 199 | return x 200 | 201 | def remove_weight_norm(self): 202 | for l in self.convs1: 203 | remove_weight_norm(l) 204 | for l in self.convs2: 205 | remove_weight_norm(l) 206 | 207 | 208 | class ResBlock2(torch.nn.Module): 209 | def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)): 210 | super(ResBlock2, self).__init__() 211 | self.h = h 212 | self.convs = nn.ModuleList([ 213 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], 214 | padding=get_padding(kernel_size, dilation[0]))), 215 | weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], 216 | padding=get_padding(kernel_size, dilation[1]))) 217 | ]) 218 | self.convs.apply(init_weights) 219 | 220 | def forward(self, x): 221 | for c in self.convs: 222 | xt = F.leaky_relu(x, LRELU_SLOPE) 223 | xt = c(xt) 224 | x = xt + x 225 | return x 226 | 227 | def remove_weight_norm(self): 228 | for l in self.convs: 229 | remove_weight_norm(l) 230 | 231 | 232 | 233 | class Generator(torch.nn.Module): 234 | def __init__(self, h): 235 | super(Generator, self).__init__() 236 | self.h = h 237 | self.num_kernels = len(h.resblock_kernel_sizes) 238 | self.num_upsamples = len(h.upsample_rates) 239 | # self.conv_pre = weight_norm(Conv1d(80, h.upsample_initial_channel, 7, 1, padding=3)) 240 | self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) # change: 80 --> 512 241 | resblock = ResBlock1 if h.resblock == '1' else ResBlock2 242 | 243 | self.ups = nn.ModuleList() 244 | for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): 245 | if (k-u) % 2 == 0: 246 | self.ups.append(weight_norm( 247 | ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)), 248 | k, u, padding=(k-u)//2))) 249 | else: 250 | self.ups.append(weight_norm( 251 | ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)), 252 | k, u, padding=(k-u)//2+1, output_padding=1))) 253 | 254 | self.resblocks = nn.ModuleList() 255 | for i in range(len(self.ups)): 256 | ch = h.upsample_initial_channel//(2**(i+1)) 257 | for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)): 258 | self.resblocks.append(resblock(h, ch, k, d)) 259 | 260 | self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3)) 261 | self.ups.apply(init_weights) 262 | self.conv_post.apply(init_weights) 263 | 264 | def forward(self, x): 265 | x = self.conv_pre(x) 266 | for i in range(self.num_upsamples): 267 | x = F.leaky_relu(x, LRELU_SLOPE) 268 | x = self.ups[i](x) 269 | xs = None 270 | for j in range(self.num_kernels): 271 | if xs is None: 272 | xs = self.resblocks[i*self.num_kernels+j](x) 273 | else: 274 | xs += self.resblocks[i*self.num_kernels+j](x) 275 | x = xs / self.num_kernels 276 | x = F.leaky_relu(x) 277 | x = self.conv_post(x) 278 | x = torch.tanh(x) 279 | 280 | return x 281 | 282 | def remove_weight_norm(self): 283 | print('Removing weight norm...') 284 | for l in self.ups: 285 | remove_weight_norm(l) 286 | for l in self.resblocks: 287 | l.remove_weight_norm() 288 | remove_weight_norm(self.conv_pre) 289 | remove_weight_norm(self.conv_post) 290 | 291 | @classmethod 292 | def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None): 293 | if subfolder is not None: 294 | pretrained_model_name_or_path = os.path.join(pretrained_model_name_or_path, subfolder) 295 | config_path = os.path.join(pretrained_model_name_or_path, "config.json") 296 | ckpt_path = os.path.join(pretrained_model_name_or_path, "vocoder.pt") 297 | 298 | config = get_config(config_path) 299 | vocoder = cls(config) 300 | 301 | state_dict_g = torch.load(ckpt_path) 302 | vocoder.load_state_dict(state_dict_g["generator"]) 303 | vocoder.eval() 304 | vocoder.remove_weight_norm() 305 | return vocoder 306 | 307 | @torch.no_grad() 308 | def inference(self, mels, lengths=None): 309 | self.eval() 310 | with torch.no_grad(): 311 | wavs = self(mels).squeeze(1) 312 | wavs = (wavs.cpu().numpy() * MAX_WAV_VALUE).astype("int16") 313 | if lengths is not None: 314 | wavs = wavs[:, :lengths] 315 | return wavs 316 | 317 | 318 | def normalize_spectrogram( 319 | spectrogram: torch.Tensor, 320 | max_value: float = 200, 321 | min_value: float = 1e-5, 322 | power: float = 1., 323 | ) -> torch.Tensor: 324 | 325 | # Rescale to 0-1 326 | max_value = np.log(max_value) # 5.298317366548036 327 | min_value = np.log(min_value) # -11.512925464970229 328 | spectrogram = torch.clamp(spectrogram, min=min_value, max=max_value) 329 | data = (spectrogram - min_value) / (max_value - min_value) 330 | # Apply the power curve 331 | data = torch.pow(data, power) 332 | # 1D -> 3D 333 | data = data.repeat(3, 1, 1) 334 | # Flip Y axis: image origin at the top-left corner, spectrogram origin at the bottom-left corner 335 | data = torch.flip(data, [1]) 336 | 337 | return data 338 | 339 | 340 | def denormalize_spectrogram( 341 | data: torch.Tensor, 342 | max_value: float = 200, 343 | min_value: float = 1e-5, 344 | power: float = 1, 345 | ) -> torch.Tensor: 346 | 347 | assert len(data.shape) == 3, "Expected 3 dimensions, got {}".format(len(data.shape)) 348 | 349 | max_value = np.log(max_value) 350 | min_value = np.log(min_value) 351 | # Flip Y axis: image origin at the top-left corner, spectrogram origin at the bottom-left corner 352 | data = torch.flip(data, [1]) 353 | if data.shape[0] == 1: 354 | data = data.repeat(3, 1, 1) 355 | assert data.shape[0] == 3, "Expected 3 channels, got {}".format(data.shape[0]) 356 | data = data[0] 357 | # Reverse the power curve 358 | data = torch.pow(data, 1 / power) 359 | # Rescale to max value 360 | spectrogram = data * (max_value - min_value) + min_value 361 | 362 | return spectrogram 363 | 364 | 365 | def pt_to_numpy(images: torch.FloatTensor) -> np.ndarray: 366 | """ 367 | Convert a PyTorch tensor to a NumPy image. 368 | """ 369 | images = images.cpu().permute(0, 2, 3, 1).float().numpy() 370 | return images 371 | 372 | 373 | def numpy_to_pil(images: np.ndarray) -> PIL.Image.Image: 374 | """ 375 | Convert a numpy image or a batch of images to a PIL image. 376 | """ 377 | if images.ndim == 3: 378 | images = images[None, ...] 379 | images = (images * 255).round().astype("uint8") 380 | if images.shape[-1] == 1: 381 | # special case for grayscale (single channel) images 382 | pil_images = [PIL.Image.fromarray(image.squeeze(), mode="L") for image in images] 383 | else: 384 | pil_images = [PIL.Image.fromarray(image) for image in images] 385 | 386 | return pil_images 387 | 388 | 389 | def image_add_color(spec_img): 390 | cmap = plt.get_cmap('viridis') 391 | cmap_r = cmap.reversed() 392 | image = cmap(np.array(spec_img)[:,:,0])[:, :, :3] # 省略透明度通道 393 | image = (image - image.min()) / (image.max() - image.min()) 394 | image = PIL.Image.fromarray(np.uint8(image*255)) 395 | return image 396 | 397 | 398 | @dataclass 399 | class PipelineOutput(BaseOutput): 400 | """ 401 | Output class for audio pipelines. 402 | 403 | Args: 404 | audios (`np.ndarray`) 405 | List of denoised audio samples of a NumPy array of shape `(batch_size, num_channels, sample_rate)`. 406 | """ 407 | 408 | images: Union[List[PIL.Image.Image], np.ndarray] 409 | spectrograms: Union[List[np.ndarray], np.ndarray] 410 | audios: Union[List[np.ndarray], np.ndarray] 411 | 412 | 413 | class AuffusionPipeline(DiffusionPipeline): 414 | 415 | r""" 416 | Pipeline for text-to-image generation using Stable Diffusion. 417 | 418 | This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the 419 | library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) 420 | 421 | In addition the pipeline inherits the following loading methods: 422 | - *Textual-Inversion*: [`loaders.TextualInversionLoaderMixin.load_textual_inversion`] 423 | - *LoRA*: [`loaders.LoraLoaderMixin.load_lora_weights`] 424 | - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`] 425 | 426 | as well as the following saving methods: 427 | - *LoRA*: [`loaders.LoraLoaderMixin.save_lora_weights`] 428 | 429 | Args: 430 | vae ([`AutoencoderKL`]): 431 | Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. 432 | text_encoder ([`CLIPTextModel`]): 433 | Frozen text-encoder. Stable Diffusion uses the text portion of 434 | [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically 435 | the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. 436 | tokenizer (`CLIPTokenizer`): 437 | Tokenizer of class 438 | [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). 439 | unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. 440 | scheduler ([`SchedulerMixin`]): 441 | A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of 442 | [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. 443 | safety_checker ([`StableDiffusionSafetyChecker`]): 444 | Classification module that estimates whether generated images could be considered offensive or harmful. 445 | Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. 446 | feature_extractor ([`CLIPImageProcessor`]): 447 | Model that extracts features from generated images to be used as inputs for the `safety_checker`. 448 | """ 449 | _optional_components = ["safety_checker", "feature_extractor", "text_encoder_list", "tokenizer_list", "adapter_list", "vocoder"] 450 | 451 | def __init__( 452 | self, 453 | vae: AutoencoderKL, 454 | unet: UNet2DConditionModel, 455 | scheduler: KarrasDiffusionSchedulers, 456 | safety_checker: StableDiffusionSafetyChecker, 457 | feature_extractor: CLIPImageProcessor, 458 | text_encoder_list: Optional[List[Callable]] = None, 459 | tokenizer_list: Optional[List[Callable]] = None, 460 | vocoder: Generator = None, 461 | requires_safety_checker: bool = False, 462 | adapter_list: Optional[List[Callable]] = None, 463 | tokenizer_model_max_length: Optional[int] = 77, # 77 is the default value for the CLIPTokenizer(and set for other models) 464 | ): 465 | super().__init__() 466 | 467 | self.text_encoder_list = text_encoder_list 468 | self.tokenizer_list = tokenizer_list 469 | self.vocoder = vocoder 470 | self.adapter_list = adapter_list 471 | self.tokenizer_model_max_length = tokenizer_model_max_length 472 | 473 | self.register_modules( 474 | vae=vae, 475 | unet=unet, 476 | scheduler=scheduler, 477 | safety_checker=safety_checker, 478 | feature_extractor=feature_extractor, 479 | ) 480 | 481 | self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) 482 | self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) 483 | self.register_to_config(requires_safety_checker=requires_safety_checker) 484 | 485 | 486 | @classmethod 487 | def from_pretrained( 488 | cls, 489 | pretrained_model_name_or_path: str = "auffusion/auffusion", 490 | dtype: torch.dtype = torch.float16, 491 | device: str = "cuda", 492 | ): 493 | if not os.path.isdir(pretrained_model_name_or_path): 494 | pretrained_model_name_or_path = snapshot_download(pretrained_model_name_or_path) 495 | 496 | vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae") 497 | unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet") 498 | feature_extractor = CLIPImageProcessor.from_pretrained(pretrained_model_name_or_path, subfolder="feature_extractor") 499 | scheduler = PNDMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler") 500 | 501 | vocoder = Generator.from_pretrained(pretrained_model_name_or_path, subfolder="vocoder").to(device, dtype) 502 | 503 | text_encoder_list, tokenizer_list, adapter_list = [], [], [] 504 | 505 | condition_json_path = os.path.join(pretrained_model_name_or_path, "condition_config.json") 506 | condition_json_list = json.loads(open(condition_json_path).read()) 507 | for i, condition_item in enumerate(condition_json_list): 508 | 509 | # Load Condition Adapter 510 | text_encoder_path = os.path.join(pretrained_model_name_or_path, condition_item["text_encoder_name"]) 511 | tokenizer = AutoTokenizer.from_pretrained(text_encoder_path) 512 | # tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder=condition_item["text_encoder_name"]) 513 | tokenizer_list.append(tokenizer) 514 | text_encoder_cls = import_model_class_from_model_name_or_path(text_encoder_path) 515 | text_encoder = text_encoder_cls.from_pretrained(text_encoder_path).to(device, dtype) 516 | text_encoder_list.append(text_encoder) 517 | print(f"LOADING CONDITION ENCODER {i}") 518 | 519 | # Load Condition Adapter 520 | adapter_path = os.path.join(pretrained_model_name_or_path, condition_item["condition_adapter_name"]) 521 | adapter = ConditionAdapter.from_pretrained(adapter_path).to(device, dtype) 522 | # adapter = ConditionAdapter.from_pretrained(pretrained_model_name_or_path, subfolder=condition_item["condition_adapter_name"]).to(device, dtype) 523 | adapter_list.append(adapter) 524 | print(f"LOADING CONDITION ADAPTER {i}") 525 | 526 | 527 | pipeline = cls( 528 | vae=vae, 529 | unet=unet, 530 | text_encoder_list=text_encoder_list, 531 | tokenizer_list=tokenizer_list, 532 | vocoder=vocoder, 533 | adapter_list=adapter_list, 534 | scheduler=scheduler, 535 | safety_checker=None, 536 | feature_extractor=feature_extractor, 537 | ) 538 | pipeline = pipeline.to(device, dtype) 539 | 540 | return pipeline 541 | 542 | 543 | def to(self, device, dtype=None): 544 | super().to(device, dtype) 545 | 546 | self.vocoder.to(device, dtype) 547 | 548 | for text_encoder in self.text_encoder_list: 549 | text_encoder.to(device, dtype) 550 | 551 | if self.adapter_list is not None: 552 | for adapter in self.adapter_list: 553 | adapter.to(device, dtype) 554 | 555 | return self 556 | 557 | 558 | def enable_vae_slicing(self): 559 | r""" 560 | Enable sliced VAE decoding. 561 | 562 | When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several 563 | steps. This is useful to save some memory and allow larger batch sizes. 564 | """ 565 | self.vae.enable_slicing() 566 | 567 | def disable_vae_slicing(self): 568 | r""" 569 | Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to 570 | computing decoding in one step. 571 | """ 572 | self.vae.disable_slicing() 573 | 574 | def enable_vae_tiling(self): 575 | r""" 576 | Enable tiled VAE decoding. 577 | 578 | When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in 579 | several steps. This is useful to save a large amount of memory and to allow the processing of larger images. 580 | """ 581 | self.vae.enable_tiling() 582 | 583 | def disable_vae_tiling(self): 584 | r""" 585 | Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to 586 | computing decoding in one step. 587 | """ 588 | self.vae.disable_tiling() 589 | 590 | def enable_sequential_cpu_offload(self, gpu_id=0): 591 | r""" 592 | Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, 593 | text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a 594 | `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called. 595 | Note that offloading happens on a submodule basis. Memory savings are higher than with 596 | `enable_model_cpu_offload`, but performance is lower. 597 | """ 598 | if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"): 599 | from accelerate import cpu_offload 600 | else: 601 | raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher") 602 | 603 | device = torch.device(f"cuda:{gpu_id}") 604 | 605 | if self.device.type != "cpu": 606 | self.to("cpu", silence_dtype_warnings=True) 607 | torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist) 608 | 609 | for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]: 610 | cpu_offload(cpu_offloaded_model, device) 611 | 612 | if self.safety_checker is not None: 613 | cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True) 614 | 615 | def enable_model_cpu_offload(self, gpu_id=0): 616 | r""" 617 | Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared 618 | to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` 619 | method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with 620 | `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. 621 | """ 622 | if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): 623 | from accelerate import cpu_offload_with_hook 624 | else: 625 | raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.") 626 | 627 | device = torch.device(f"cuda:{gpu_id}") 628 | 629 | if self.device.type != "cpu": 630 | self.to("cpu", silence_dtype_warnings=True) 631 | torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist) 632 | 633 | hook = None 634 | for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]: 635 | _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook) 636 | 637 | if self.safety_checker is not None: 638 | _, hook = cpu_offload_with_hook(self.safety_checker, device, prev_module_hook=hook) 639 | 640 | # We'll offload the last model manually. 641 | self.final_offload_hook = hook 642 | 643 | @property 644 | def _execution_device(self): 645 | r""" 646 | Returns the device on which the pipeline's models will be executed. After calling 647 | `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module 648 | hooks. 649 | """ 650 | if not hasattr(self.unet, "_hf_hook"): 651 | return self.device 652 | for module in self.unet.modules(): 653 | if ( 654 | hasattr(module, "_hf_hook") 655 | and hasattr(module._hf_hook, "execution_device") 656 | and module._hf_hook.execution_device is not None 657 | ): 658 | return torch.device(module._hf_hook.execution_device) 659 | return self.device 660 | 661 | 662 | def _encode_prompt( 663 | self, 664 | prompt, 665 | device, 666 | num_images_per_prompt, 667 | do_classifier_free_guidance, 668 | negative_prompt=None, 669 | prompt_embeds: Optional[torch.FloatTensor] = None, 670 | negative_prompt_embeds: Optional[torch.FloatTensor] = None, 671 | ): 672 | 673 | assert len(self.text_encoder_list) == len(self.tokenizer_list), "Number of text_encoders must match number of tokenizers" 674 | if self.adapter_list is not None: 675 | assert len(self.text_encoder_list) == len(self.adapter_list), "Number of text_encoders must match number of adapters" 676 | 677 | if prompt is not None and isinstance(prompt, str): 678 | batch_size = 1 679 | elif prompt is not None and isinstance(prompt, list): 680 | batch_size = len(prompt) 681 | else: 682 | batch_size = prompt_embeds.shape[0] 683 | 684 | def get_prompt_embeds(prompt_list, device): 685 | if isinstance(prompt_list, str): 686 | prompt_list = [prompt_list] 687 | 688 | prompt_embeds_list = [] 689 | for prompt in prompt_list: 690 | encoder_hidden_states_list = [] 691 | 692 | # Generate condition embedding 693 | for j in range(len(self.text_encoder_list)): 694 | # get condition embedding using condition encoder 695 | input_ids = self.tokenizer_list[j](prompt, return_tensors="pt").input_ids.to(device) 696 | cond_embs = self.text_encoder_list[j](input_ids).last_hidden_state # [bz, text_len, text_dim] 697 | # padding to max_length 698 | if cond_embs.shape[1] < self.tokenizer_model_max_length: 699 | cond_embs = torch.functional.F.pad(cond_embs, (0, 0, 0, self.tokenizer_model_max_length - cond_embs.shape[1]), value=0) 700 | else: 701 | cond_embs = cond_embs[:, :self.tokenizer_model_max_length, :] 702 | 703 | # use condition adapter 704 | if self.adapter_list is not None: 705 | cond_embs = self.adapter_list[j](cond_embs) 706 | encoder_hidden_states_list.append(cond_embs) 707 | 708 | prompt_embeds = torch.cat(encoder_hidden_states_list, dim=1) 709 | prompt_embeds_list.append(prompt_embeds) 710 | 711 | prompt_embeds = torch.cat(prompt_embeds_list, dim=0) 712 | return prompt_embeds 713 | 714 | 715 | if prompt_embeds is None: 716 | prompt_embeds = get_prompt_embeds(prompt, device) 717 | 718 | prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device) 719 | 720 | bs_embed, seq_len, _ = prompt_embeds.shape 721 | # duplicate text embeddings for each generation per prompt, using mps friendly method 722 | prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) 723 | prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) 724 | 725 | if do_classifier_free_guidance and negative_prompt_embeds is None: 726 | 727 | if negative_prompt is None: 728 | negative_prompt_embeds = torch.zeros_like(prompt_embeds).to(dtype=prompt_embeds.dtype, device=device) 729 | 730 | elif prompt is not None and type(prompt) is not type(negative_prompt): 731 | raise TypeError( 732 | f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" 733 | f" {type(prompt)}." 734 | ) 735 | elif isinstance(negative_prompt, str): 736 | negative_prompt = [negative_prompt] 737 | negative_prompt_embeds = get_prompt_embeds(negative_prompt, device) 738 | elif batch_size != len(negative_prompt): 739 | raise ValueError( 740 | f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" 741 | f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" 742 | " the batch size of `prompt`." 743 | ) 744 | else: 745 | negative_prompt_embeds = get_prompt_embeds(negative_prompt, device) 746 | 747 | if do_classifier_free_guidance: 748 | # duplicate unconditional embeddings for each generation per prompt, using mps friendly method 749 | seq_len = negative_prompt_embeds.shape[1] 750 | 751 | negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device) 752 | 753 | negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) 754 | negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) 755 | 756 | # For classifier free guidance, we need to do two forward passes. 757 | # Here we concatenate the unconditional and text embeddings into a single batch 758 | # to avoid doing two forward passes 759 | prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) 760 | 761 | return prompt_embeds 762 | 763 | 764 | def run_safety_checker(self, image, device, dtype): 765 | if self.safety_checker is None: 766 | has_nsfw_concept = None 767 | else: 768 | if torch.is_tensor(image): 769 | feature_extractor_input = self.image_processor.postprocess(image, output_type="pil") 770 | else: 771 | feature_extractor_input = self.image_processor.numpy_to_pil(image) 772 | safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device) 773 | image, has_nsfw_concept = self.safety_checker( 774 | images=image, clip_input=safety_checker_input.pixel_values.to(dtype) 775 | ) 776 | return image, has_nsfw_concept 777 | 778 | def decode_latents(self, latents): 779 | warnings.warn( 780 | "The decode_latents method is deprecated and will be removed in a future version. Please" 781 | " use VaeImageProcessor instead", 782 | FutureWarning, 783 | ) 784 | latents = 1 / self.vae.config.scaling_factor * latents 785 | image = self.vae.decode(latents, return_dict=False)[0] 786 | image = (image / 2 + 0.5).clamp(0, 1) 787 | # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 788 | image = image.cpu().permute(0, 2, 3, 1).float().numpy() 789 | return image 790 | 791 | def prepare_extra_step_kwargs(self, generator, eta): 792 | # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature 793 | # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. 794 | # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 795 | # and should be between [0, 1] 796 | 797 | accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) 798 | extra_step_kwargs = {} 799 | if accepts_eta: 800 | extra_step_kwargs["eta"] = eta 801 | 802 | # check if the scheduler accepts generator 803 | accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) 804 | if accepts_generator: 805 | extra_step_kwargs["generator"] = generator 806 | return extra_step_kwargs 807 | 808 | 809 | def check_inputs( 810 | self, 811 | prompt, 812 | height, 813 | width, 814 | callback_steps, 815 | negative_prompt=None, 816 | prompt_embeds=None, 817 | negative_prompt_embeds=None, 818 | ): 819 | if height % 8 != 0 or width % 8 != 0: 820 | raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") 821 | 822 | if (callback_steps is None) or ( 823 | callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) 824 | ): 825 | raise ValueError( 826 | f"`callback_steps` has to be a positive integer but is {callback_steps} of type" 827 | f" {type(callback_steps)}." 828 | ) 829 | 830 | if prompt is not None and prompt_embeds is not None: 831 | raise ValueError( 832 | f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" 833 | " only forward one of the two." 834 | ) 835 | elif prompt is None and prompt_embeds is None: 836 | raise ValueError( 837 | "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." 838 | ) 839 | elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): 840 | raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") 841 | 842 | if negative_prompt is not None and negative_prompt_embeds is not None: 843 | raise ValueError( 844 | f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" 845 | f" {negative_prompt_embeds}. Please make sure to only forward one of the two." 846 | ) 847 | 848 | if prompt_embeds is not None and negative_prompt_embeds is not None: 849 | if prompt_embeds.shape != negative_prompt_embeds.shape: 850 | raise ValueError( 851 | "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" 852 | f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" 853 | f" {negative_prompt_embeds.shape}." 854 | ) 855 | 856 | 857 | 858 | def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): 859 | shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) 860 | if isinstance(generator, list) and len(generator) != batch_size: 861 | raise ValueError( 862 | f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" 863 | f" size of {batch_size}. Make sure the batch size matches the length of the generators." 864 | ) 865 | 866 | if latents is None: 867 | latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) 868 | else: 869 | latents = latents.to(device) 870 | 871 | # scale the initial noise by the standard deviation required by the scheduler 872 | latents = latents * self.scheduler.init_noise_sigma 873 | return latents 874 | 875 | def aggregate_attention(self, prompts, attention_store: AttentionStore, res: List[int], from_where: List[str], is_cross: bool, select: int): 876 | out = [] 877 | attention_maps = attention_store.get_average_attention() 878 | # num_pixels = res ** 2 879 | num_pixels = res[0] * res[1] 880 | for location in from_where: 881 | for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]: 882 | if item.shape[1] == num_pixels: 883 | cross_maps = item.reshape(len(prompts), -1, res[0], res[1], item.shape[-1])[select] 884 | out.append(cross_maps) 885 | out = torch.cat(out, dim=0) 886 | out = out.sum(0) / out.shape[0] 887 | return out.cpu() 888 | 889 | @torch.no_grad() 890 | def __call__( 891 | self, 892 | prompt: Union[str, List[str]] = None, 893 | height: Optional[int] = 256, 894 | width: Optional[int] = 1024, 895 | num_inference_steps: int = 100, 896 | guidance_scale: float = 7.5, 897 | negative_prompt: Optional[Union[str, List[str]]] = None, 898 | num_images_per_prompt: Optional[int] = 1, 899 | eta: float = 0.0, 900 | generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, 901 | latents: Optional[torch.FloatTensor] = None, 902 | prompt_embeds: Optional[torch.FloatTensor] = None, 903 | negative_prompt_embeds: Optional[torch.FloatTensor] = None, 904 | output_type: Optional[str] = "pt", 905 | return_dict: bool = True, 906 | callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, 907 | callback_steps: int = 1, 908 | cross_attention_kwargs: Optional[Dict[str, Any]] = None, 909 | guidance_rescale: float = 0.0, 910 | duration: Optional[float] = 10, 911 | ): 912 | 913 | # 0. Default height and width to unet 914 | height = height or self.unet.config.sample_size * self.vae_scale_factor 915 | width = width or self.unet.config.sample_size * self.vae_scale_factor 916 | # print(height) # 256 917 | # print(width) # 1024 918 | audio_length = int(duration * 16000) 919 | 920 | 921 | # 1. Check inputs. Raise error if not correct 922 | self.check_inputs( 923 | prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds 924 | ) 925 | 926 | 927 | # 2. Define call parameters 928 | if prompt is not None and isinstance(prompt, str): 929 | batch_size = 1 930 | elif prompt is not None and isinstance(prompt, list): 931 | batch_size = len(prompt) 932 | else: 933 | batch_size = prompt_embeds.shape[0] 934 | 935 | device = self._execution_device 936 | # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) 937 | # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` 938 | # corresponds to doing no classifier free guidance. 939 | do_classifier_free_guidance = guidance_scale > 1.0 940 | 941 | # 3. Encode input prompt 942 | prompt_embeds = self._encode_prompt( 943 | prompt, 944 | device, 945 | num_images_per_prompt, 946 | do_classifier_free_guidance, 947 | negative_prompt, 948 | prompt_embeds=prompt_embeds, 949 | negative_prompt_embeds=negative_prompt_embeds 950 | ) 951 | 952 | # print(prompt_embeds.shape) # torch.Size([2, 77, 768]) 953 | # print(prompt_embeds) 954 | 955 | # 4. Prepare timesteps 956 | self.scheduler.set_timesteps(num_inference_steps, device=device) 957 | timesteps = self.scheduler.timesteps 958 | 959 | # 5. Prepare latent variables 960 | num_channels_latents = self.unet.config.in_channels 961 | latents = self.prepare_latents( 962 | batch_size * num_images_per_prompt, 963 | num_channels_latents, 964 | height, 965 | width, 966 | prompt_embeds.dtype, 967 | device, 968 | generator, 969 | latents, 970 | ) 971 | 972 | # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline 973 | extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) 974 | 975 | # 7. Denoising loop 976 | num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order 977 | with self.progress_bar(total=num_inference_steps) as progress_bar: 978 | for i, t in enumerate(timesteps): 979 | # expand the latents if we are doing classifier free guidance 980 | latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents 981 | latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) 982 | 983 | # predict the noise residual 984 | noise_pred = self.unet( 985 | latent_model_input, 986 | t, 987 | encoder_hidden_states=prompt_embeds, 988 | cross_attention_kwargs=cross_attention_kwargs, 989 | return_dict=False, 990 | )[0] 991 | 992 | # perform guidance 993 | if do_classifier_free_guidance: 994 | noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) 995 | noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) 996 | 997 | if do_classifier_free_guidance and guidance_rescale > 0.0: 998 | # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf 999 | noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) 1000 | 1001 | # compute the previous noisy sample x_t -> x_t-1 1002 | latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] 1003 | 1004 | # call the callback, if provided 1005 | if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): 1006 | progress_bar.update() 1007 | if callback is not None and i % callback_steps == 0: 1008 | callback(i, t, latents) 1009 | 1010 | if not output_type == "latent": 1011 | image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] 1012 | image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) 1013 | else: 1014 | image = latents 1015 | has_nsfw_concept = None 1016 | 1017 | if has_nsfw_concept is None: 1018 | do_denormalize = [True] * image.shape[0] 1019 | else: 1020 | do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] 1021 | 1022 | image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) 1023 | 1024 | # Offload last model to CPU 1025 | if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: 1026 | self.final_offload_hook.offload() 1027 | 1028 | 1029 | # Generate audio 1030 | spectrograms, audios = [], [] 1031 | for img in image: 1032 | spectrogram = denormalize_spectrogram(img) 1033 | audio = self.vocoder.inference(spectrogram, lengths=audio_length)[0] 1034 | audios.append(audio) 1035 | spectrograms.append(spectrogram) 1036 | 1037 | # Convert to PIL 1038 | images = pt_to_numpy(image) 1039 | images = numpy_to_pil(images) 1040 | images = [image_add_color(image) for image in images] 1041 | 1042 | if not return_dict: 1043 | return (images, audios, spectrograms) 1044 | 1045 | 1046 | return PipelineOutput(images=images, audios=audios, spectrograms=spectrograms) 1047 | 1048 | 1049 | 1050 | 1051 | --------------------------------------------------------------------------------