Upload 3 files

Add Inference Files

app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+import os
+import sys
+sys.path.insert(1, '..')
+import inference
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def generate_background_music(video_path, dataset, use_peft, musicgen_size):
+    print(f"Start generating background music for {video_path} with model \"{'peft' if use_peft else 'audiocraft'}_{dataset}_{musicgen_size}\"")
+
+    new_video_path = inference.generate_background_music(
+        video_path=video_path,
+        dataset=dataset,
+        musicgen_size=musicgen_size,
+        use_stereo=True,
+        use_peft=use_peft,
+        musicgen_temperature=1.0,
+        musicgen_guidance_scale=3.0,
+        top_k_sampling=250,
+        device=device
+    )
+    return gr.Video(new_video_path)
+
+
+interface = gr.Interface(fn=generate_background_music,
+                         inputs=[
+                                 gr.Video(
+                                          label="video input",
+                                          min_length=5,
+                                          max_length=20,
+                                          sources=['upload'],
+                                          show_download_button=True,
+                                          include_audio=True
+                                          ),
+                                 gr.Radio(["nature", "symmv"],
+                                          label="Video Encoder Version",
+                                          value="nature",
+                                          info="Choose one of the available Video Encoders."),
+                                 gr.Radio([False, True],
+                                          label="Use MusicGen Audio Decoder Model trained with PEFT",
+                                          value=False,
+                                          info="If set to 'True' the MusicGen Audio Decoder models trained with LoRA "
+                                               "(Low Rank Adaptation) are used. If set to 'False', the original "
+                                               "MusicGen models are used."),
+                                 gr.Radio(["small", "medium", "large"],
+                                          label="MusicGen Audio Decoder Size",
+                                          value="small",
+                                          info="Choose the size of the MusicGen audio decoder."),
+                                ],
+
+                         outputs=[gr.Video(label="video output")],
+                         examples=[
+                             [os.path.abspath("../../../videos/originals/n_1.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_2.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_3.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_4.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_5.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_6.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_7.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_8.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_1.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_2.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_3.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_4.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_5.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_6.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_7.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_8.mp4"), "nature", True, "small"],
+                           ],
+                         cache_examples=False
+                         )
+
+if __name__ == "__main__":
+    interface.launch(
+        share=False
+    )

inference.py

@@ -0,0 +1,117 @@
+from omegaconf import OmegaConf
+from peft import PeftConfig, get_peft_model 
+from audiocraft.models import MusicGen
+from moviepy.editor import AudioFileClip
+from code.inference.inference_utils import *
+import re
+import time
+
+re_file_name = re.compile('([^/]+$)')
+
+
+def generate_background_music(video_path: str,
+                              dataset: str,
+                              musicgen_size: str,
+                              use_stereo: bool,
+                              use_peft: bool,
+                              device: str,
+                              musicgen_temperature: float = 1.0,
+                              musicgen_guidance_scale: float = 3.0,
+                              top_k_sampling: int = 250) -> str:
+    start = time.time()
+    model_path = "../training/"
+    model_path += "models_peft" if use_peft else "models_audiocraft"
+    model_path += f"/{dataset}" + f"_{musicgen_size}"
+
+    conf = OmegaConf.load(model_path + '/configuration.yml')
+    use_sampling = True if top_k_sampling > 0 else False
+    video = mpe.VideoFileClip(video_path)
+
+    musicgen_model_id = "facebook/musicgen-" + "stereo-" if use_stereo else ""
+    musicgen_model_id += musicgen_size
+
+    result_dir = "./results"
+    os.makedirs(result_dir, exist_ok=True)
+
+    encoder_output_dimension = None
+    if "small" in conf.musicgen_model_id:
+        encoder_output_dimension = 1024
+    elif "medium" in conf.musicgen_model_id:
+        encoder_output_dimension = 1536
+    elif "large" in conf.musicgen_model_id:
+        encoder_output_dimension = 2048
+    assert encoder_output_dimension, f"Video Encoder output dimension could not be determined by {conf.musicgen_model_id}"
+
+    musicgen_model = MusicGen.get_pretrained(musicgen_model_id)
+    musicgen_model.lm.to(device)
+    musicgen_model.compression_model.to(device)
+    if use_peft:
+        peft_path = model_path + "/musicgen_peft_final"
+        peft_config = PeftConfig.from_pretrained(peft_path)
+        musicgen_model.lm = get_peft_model(musicgen_model.lm, peft_config)
+        musicgen_model.lm.load_adapter(peft_path, "default")
+
+    print("MusicGen Model loaded.")
+
+    video_to_t5 = VideoToT5(
+        video_extraction_framerate=conf.video_extraction_framerate,
+        encoder_input_dimension=conf.encoder_input_dimension,
+        encoder_output_dimension=encoder_output_dimension,
+        encoder_heads=conf.encoder_heads,
+        encoder_dim_feedforward=conf.encoder_dim_feedforward,
+        encoder_layers=conf.encoder_layers,
+        device=device
+    )
+
+    video_to_t5.load_state_dict(torch.load(model_path + "/lm_final.pt", map_location=device))
+    print("Video Encoder Model loaded.")
+
+    print("Starting Video Feature Extraction.")
+    video_embedding_t5 = video_to_t5(video_paths=[video_path])
+
+    condition_tensors = create_condition_tensors(
+        video_embeddings=video_embedding_t5,
+        batch_size=1,
+        video_extraction_framerate=video_to_t5.video_extraction_framerate,
+        device=device
+    )
+
+    musicgen_model.generation_params = {
+        'max_gen_len': int(video.duration * musicgen_model.frame_rate),
+        'use_sampling': use_sampling,
+        'temp': musicgen_temperature,
+        'cfg_coef': musicgen_guidance_scale,
+        'two_step_cfg': False,
+    }
+    if use_sampling:
+        musicgen_model.generation_params['top_k'] = 250
+
+    print("Starting Audio Generation.")
+    prompt_tokens = None
+    with torch.no_grad():
+        with musicgen_model.autocast:
+            gen_tokens = musicgen_model.lm.generate(prompt_tokens, [], condition_tensors, callback=None,
+                                                    **musicgen_model.generation_params)
+        gen_audio = musicgen_model.compression_model.decode(gen_tokens)
+
+    end = time.time()
+    print("Elapsed time for generation: " + str(end - start))
+
+    _, video_file_name = os.path.split(video_path)
+    video_file_name = video_file_name[:-4]  # remove .mp4
+
+    re_result = re_file_name.search(video_file_name)  # get video file name
+    result_path = f"{'peft' if use_peft else 'audiocraft'}_{dataset}_{musicgen_size}_{re_result.group(1)}"
+    audio_result_path = f"{result_dir}/tmp.wav"
+    video_result_path = f"{result_dir}/{result_path}_video.mp4"
+
+    gen_audio = torch.squeeze(gen_audio.detach().cpu())  # remove mini-batch dimension, move to CPU for saving
+    sample_rate = musicgen_model.sample_rate
+    torchaudio.save(audio_result_path, gen_audio, sample_rate)
+    audio_file_clip = AudioFileClip(audio_result_path)
+    video.audio = audio_file_clip
+
+    print("Rendering Video.")
+    video.write_videofile(video_result_path)
+
+    return video_result_path

inference_utils.py

@@ -0,0 +1,283 @@
+from torch.utils.data import Dataset
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+import torchaudio
+import os
+import logging
+from torchvision.models import resnet50, ResNet50_Weights, resnet152, resnet18, resnet34, ResNet152_Weights
+from PIL import Image
+from time import strftime
+import math
+import numpy as np
+import moviepy.editor as mpe
+
+
+class VideoDataset(Dataset):
+    def __init__(self, data_dir):
+        self.data_dir = data_dir
+        self.data_map = []
+
+        dir_map = os.listdir(data_dir)
+        for d in dir_map:
+            name, extension = os.path.splitext(d)
+            if extension == ".mp4":
+                self.data_map.append({"video": os.path.join(data_dir, d)})
+
+    def __len__(self):
+        return len(self.data_map)
+
+    def __getitem__(self, idx):
+        return self.data_map[idx]["video"]
+
+
+# input: video_path, output: wav_music
+class VideoToT5(nn.Module):
+    def __init__(self,
+                 device: str,
+                 video_extraction_framerate: int,
+                 encoder_input_dimension: int,
+                 encoder_output_dimension: int,
+                 encoder_heads: int,
+                 encoder_dim_feedforward: int,
+                 encoder_layers: int
+                 ):
+        super().__init__()
+        self.video_extraction_framerate = video_extraction_framerate
+        self.video_feature_extractor = VideoFeatureExtractor(video_extraction_framerate=video_extraction_framerate,
+                                                             device=device)
+        self.video_encoder = VideoEncoder(
+            device,
+            encoder_input_dimension,
+            encoder_output_dimension,
+            encoder_heads,
+            encoder_dim_feedforward,
+            encoder_layers
+        )
+
+    def forward(self, video_paths: [str]):
+        image_embeddings = []
+        for video_path in video_paths:
+            video = mpe.VideoFileClip(video_path)
+            video_embedding = self.video_feature_extractor(video)
+            image_embeddings.append(video_embedding)
+        video_embedding = torch.stack(
+            image_embeddings)  # resulting shape: [batch_size, video_extraction_framerate, resnet_output_dimension]
+        # not used, gives worse results!
+        # video_embeddings = torch.mean(video_embeddings, 0, True)  # average out all image embedding to one video embedding
+
+        t5_embeddings = self.video_encoder(video_embedding)  # T5 output: [batch_size, num_tokens,
+        # t5_embedding_size]
+        return t5_embeddings
+
+
+class VideoEncoder(nn.Module):
+    def __init__(self,
+                 device: str,
+                 encoder_input_dimension: int,
+                 encoder_output_dimension: int,
+                 encoder_heads: int,
+                 encoder_dim_feedforward: int,
+                 encoder_layers: int
+                 ):
+        super().__init__()
+        self.device = device
+        self.encoder = (nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=encoder_input_dimension,
+                nhead=encoder_heads,
+                dim_feedforward=encoder_dim_feedforward
+            ),
+            num_layers=encoder_layers,
+        )
+        ).to(device)
+
+        # linear layer to match T5 embedding dimension
+        self.linear = (nn.Linear(
+            in_features=encoder_input_dimension,
+            out_features=encoder_output_dimension)
+                       .to(device))
+
+    def forward(self, x):
+        assert x.dim() == 3
+        x = torch.transpose(x, 0, 1)  # encoder expects [sequence_length, batch_size, embedding_dimension]
+        x = self.encoder(x)  # encoder forward pass
+        x = self.linear(x)  # forward pass through the linear layer
+        x = torch.transpose(x, 0, 1)  # shape: [batch_size, sequence_length, embedding_dimension]
+        return x
+
+
+class VideoFeatureExtractor(nn.Module):
+    def __init__(self,
+                 device: str,
+                 video_extraction_framerate: int = 1,
+                 resnet_output_dimension: int = 2048):
+        super().__init__()
+        self.device = device
+
+        # using a ResNet trained on ImageNet
+        self.resnet = resnet50(weights="IMAGENET1K_V2").eval()
+        self.resnet = torch.nn.Sequential(*(list(self.resnet.children())[:-1])).to(device)  # remove ResNet layer
+        self.resnet_preprocessor = ResNet50_Weights.DEFAULT.transforms().to(device)
+        self.video_extraction_framerate = video_extraction_framerate  # setting the fps at which the video is processed
+        self.positional_encoder = PositionalEncoding(resnet_output_dimension).to(device)
+
+    def forward(self, video: mpe.VideoFileClip):
+        embeddings = []
+        for i in range(0, 30 * self.video_extraction_framerate):
+            i = video.get_frame(i)  # get frame as numpy array
+            i = Image.fromarray(i)  # create PIL image from numpy array
+            i = self.resnet_preprocessor(i)  # preprocess image
+            i = i.to(self.device)
+            i = i.unsqueeze(0)  # adding a batch dimension
+            i = self.resnet(i).squeeze()  # ResNet forward pass
+            i = i.squeeze()
+            embeddings.append(i)  # collect embeddings
+
+        embeddings = torch.stack(embeddings)  # concatenate all frame embeddings into one video embedding
+        embeddings = embeddings.unsqueeze(1)
+        embeddings = self.positional_encoder(embeddings)  # apply positional encoding with a sequence length of 30
+        embeddings = embeddings.squeeze()
+        return embeddings
+
+
+# from https://pytorch.org/tutorials/beginner/transformer_tutorial.html
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, dropout: float = 0.1, max_length: int = 5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        position = torch.arange(30).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(30, 1, d_model)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.pe[:x.size(0)]
+        return self.dropout(x)
+
+
+def freeze_model(model: nn.Module):
+    for param in model.parameters():
+        param.requires_grad = False
+    model.eval()
+
+
+def split_dataset_randomly(dataset, validation_split: float, test_split: float, seed: int = None):
+    dataset_size = len(dataset)
+    indices = list(range(dataset_size))
+    datapoints_validation = int(np.floor(validation_split * dataset_size))
+    datapoints_testing = int(np.floor(test_split * dataset_size))
+
+    if seed:
+        np.random.seed(seed)
+
+    np.random.shuffle(indices)  # in-place operation
+    training = indices[datapoints_validation + datapoints_testing:]
+    validation = indices[datapoints_validation:datapoints_testing + datapoints_validation]
+    testing = indices[:datapoints_testing]
+
+    assert len(validation) == datapoints_validation, "Validation set length incorrect"
+    assert len(testing) == datapoints_testing, "Testing set length incorrect"
+    assert len(training) == dataset_size - (datapoints_testing + datapoints_testing), "Training set length incorrect"
+    assert not any([item in training for item in validation]), "Training and Validation overlap"
+    assert not any([item in training for item in testing]), "Training and Testing overlap"
+    assert not any([item in validation for item in testing]), "Validation and Testing overlap"
+
+    return training, validation, testing
+
+
+### private function from audiocraft.solver.musicgen.py => _compute_cross_entropy
+def compute_cross_entropy(logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor):
+    """Compute cross entropy between multi-codebook targets and model's logits.
+    The cross entropy is computed per codebook to provide codebook-level cross entropy.
+    Valid timesteps for each of the codebook are pulled from the mask, where invalid
+    timesteps are set to 0.
+
+    Args:
+        logits (torch.Tensor): Model's logits of shape [B, K, T, card].
+        targets (torch.Tensor): Target codes, of shape [B, K, T].
+        mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T].
+    Returns:
+        ce (torch.Tensor): Cross entropy averaged over the codebooks
+        ce_per_codebook (list of torch.Tensor): Cross entropy per codebook (detached).
+    """
+    B, K, T = targets.shape
+    assert logits.shape[:-1] == targets.shape
+    assert mask.shape == targets.shape
+    ce = torch.zeros([], device=targets.device)
+    ce_per_codebook = []
+    for k in range(K):
+        logits_k = logits[:, k, ...].contiguous().view(-1, logits.size(-1))  # [B x T, card]
+        targets_k = targets[:, k, ...].contiguous().view(-1)  # [B x T]
+        mask_k = mask[:, k, ...].contiguous().view(-1)  # [B x T]
+        ce_targets = targets_k[mask_k]
+        ce_logits = logits_k[mask_k]
+        q_ce = F.cross_entropy(ce_logits, ce_targets)
+        ce += q_ce
+        ce_per_codebook.append(q_ce.detach())
+    # average cross entropy across codebooks
+    ce = ce / K
+    return ce, ce_per_codebook
+
+
+def generate_audio_codes(audio_paths: [str],
+                         audiocraft_compression_model: torch.nn.Module,
+                         device: str) -> torch.Tensor:
+    audio_duration = 30
+    encodec_sample_rate = audiocraft_compression_model.sample_rate
+
+    torch_audios = []
+    for audio_path in audio_paths:
+        wav, original_sample_rate = torchaudio.load(audio_path)  # load audio from file
+        wav = torchaudio.functional.resample(wav, original_sample_rate,
+                                             encodec_sample_rate)  # cast audio to model sample rate
+        wav = wav[:, :encodec_sample_rate * audio_duration]  # enforce an exact audio length of 30 seconds
+
+        assert len(wav.shape) == 2, f"audio data is not of shape [channels, duration]"
+        assert wav.shape[0] == 2, "audio data should be in stereo, but has not 2 channels"
+
+        torch_audios.append(wav)
+
+    torch_audios = torch.stack(torch_audios)
+    torch_audios = torch_audios.to(device)
+
+    with torch.no_grad():
+        gen_audio = audiocraft_compression_model.encode(torch_audios)
+
+    codes, scale = gen_audio
+    assert scale is None
+
+    return codes
+
+
+def create_condition_tensors(
+        video_embeddings: torch.Tensor,
+        batch_size: int,
+        video_extraction_framerate: int,
+        device: str
+):
+    # model T5 mask
+    mask = torch.ones((batch_size, video_extraction_framerate * 30), dtype=torch.int).to(device)
+
+    condition_tensors = {
+        'description': (video_embeddings, mask)
+    }
+    return condition_tensors
+
+
+def get_current_timestamp():
+    return strftime("%Y_%m_%d___%H_%M_%S")
+
+
+def configure_logging(output_dir: str, filename: str, log_level):
+    # create logs folder, if not existing
+    os.makedirs(output_dir, exist_ok=True)
+    level = getattr(logging, log_level)
+    file_path = output_dir + "/" + filename
+    logging.basicConfig(filename=file_path, encoding='utf-8', level=level)
+    logger = logging.getLogger()
+    # only add a StreamHandler if it is not present yet
+    if len(logger.handlers) <= 1:
+        logger.addHandler(logging.StreamHandler())