Upload 80 files

.gitattributes

@@ -33,3 +33,27 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio1/bass.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio1/drums.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio1/other.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio1/vocals.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio2/bass.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio2/drums.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio2/other.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio2/vocals.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio3/bass.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio3/drums.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio3/other.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio3/vocals.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio4/bass.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio4/drums.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio4/other.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio4/vocals.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio5/bass.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio5/drums.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio5/other.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio5/vocals.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio6/bass.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio6/drums.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio6/other.wav filter=lfs diff=lfs merge=lfs -text
+JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio6/vocals.wav filter=lfs diff=lfs merge=lfs -text

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/README.md

@@ -0,0 +1,99 @@
+# JEN-1-COMPOSER-pytorch(WIP)
+Unofficial implementation JEN-1 Composer: A Unified Framework for High-Fidelity Multi-Track Music Generation(https://arxiv.org/abs/2310.19180)
+
+![JEN-1](https://github.com/0417keito/JEN-1-COMPOSER-pytorch/blob/main/JEN1-Composer.jpg)
+![JEN-1-fig2](https://github.com/0417keito/JEN-1-COMPOSER-pytorch/blob/main/Jen1-Composer-2.png)
+
+## README
+
+## 📖 Quick Index
+* [💻 Installation](#-installation)
+* [🐍Usage](#-method)
+* [🧠TODO](#-todo)
+* [🚀Demo](#-demo)
+* [🙏Appreciation](#-appreciation)
+* [⭐️Show Your Support](#-show_your_support)
+* [🙆Welcome Contributions](#-welcom_contributions)
+
+## 💻 Installation
+```commandline
+git clone https://github.com/0417keito/JEN-1-pytorch.git
+cd JEN-1-pytorch
+pip install -r requirements.txt
+```
+
+## 🐍Usage
+### Sampling
+```python
+import torch
+from generation import Jen1
+
+ckpt_path =  'your ckpt path'
+jen1 = Jen1(ckpt_path)
+
+prompt = 'a beautiful song'
+samples = jen1.generate(prompt)
+```
+
+### Training
+```commandline
+torchrun train.py
+```
+
+### Dataset format
+Json format. the name of the Json file must be the same as the target music file.
+```json
+{"prompt": "a beautiful song"}
+```
+```python
+How should the data_dir be created?
+
+'''
+dataset_dir
+├── audios
+|    ├── music1.wav
+|    ├── music2.wav
+|    .......
+|    ├── music{n}.wav
+|
+├── metadata
+|   ├── music1.json
+|   ├── music2.json
+|   ......
+|   ├── music{n}.json
+|
+'''
+```
+
+### About config
+please see [config.py](https://github.com/0417keito/JEN-1-pytorch/blob/main/utils/config.py) and [conditioner_config.py](https://github.com/0417keito/JEN-1-pytorch/blob/main/utils/conditioner_config.py)
+
+## 🧠TODO
+- [ ] Extension to [JEN-1-Composer](https://arxiv.org/abs/2310.19180) 
+- [ ] Extension to music generation with singing voice
+- [ ] Adaptation of Consistency Model
+- [ ] In the paper, Diffusion Autoencoder was used, but I did not have much computing resources, so I used Encodec instead. So, if I can afford it, I will implement Diffusion Autoencoder.
+
+## 🚀Demo
+coming soon !
+
+## 🙏Appreciation
+[Dr Adam Fils](https://github.com/adamfils) - for support and brought this to my attention.
+
+## ⭐️Show Your Support
+
+If you find this repo interesting and useful, give us a ⭐️ on GitHub! It encourages us to keep improving the model and adding exciting features.
+Please inform us of any deficiencies by issue.
+
+## 🙆Welcome Contributions
+Contributions are always welcome.
+
+## Citations
+```bibtex
+@misc{2310.19180,
+Author = {Yao Yao and Peike Li and Boyu Chen and Alex Wang},
+Title = {JEN-1 Composer: A Unified Framework for High-Fidelity Multi-Track Music Generation},
+Year = {2023},
+Eprint = {arXiv:2310.19180},
+}
+```

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio1/bass.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f6bc7ca834b5899a1f35e0294213ab0b74a305c718639af75b5d5f079cd7ca7
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio1/drums.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e0a30f4be76598d557f5eb3e4c616f104d8b6624484de64ef040acf790d2380
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio1/other.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39447a94082b91122e123e93edaa7fae768e796d6092eb1471f67e38a30d5253
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio1/vocals.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3c51fb6ffd95e13cc866c499037f84a972985decb9cdd93c184fa27cd464458
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio2/bass.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8fb9f3d14c2f7a9c6792fe686fafc9a9be9e8169ab5efd481bd274156af2f18
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio2/drums.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aad842cce67224f86c8bd3aa8e6cc673467b3f02dc39883be66a662b3c34363d
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio2/other.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4d1bee50855e4c06252e665300a9ff02b1aa1c66822bb60da0119ec950b0299
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio2/vocals.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca8c87f0b2569e52d499601463b4f0f394bad3e0c8002acca3e0aad25bc356f8
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio3/bass.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2bf1d2592748b07e4a8553aae79eacb25b1d19c50f2ec7b658c9a224138e18d8
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio3/drums.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df8950751972beb6fc0f628a6c53e230f0eb785e68d9b3f8c8ea62895aaa40c5
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio3/other.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c003b116a0f3e56addfb7de9b3fe5e60c0a70e8d0d1f4af28b7e4bebd0639c2e
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio3/vocals.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9c62024fc680006609ee8c01c0bf8092287b7f8d435a979cf200eeb3de3b10f
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio4/bass.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93dee392485d0c3d3f9ca3cd312dae648a5840135cf272df631a0c815fbd0449
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio4/drums.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed3e0f48e874df199d8f95392504dde16d727da2396bb3427704bd7ec105c5a1
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio4/other.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1660aa65126295133d2fe1cc1ba12643c2679e13ecfa864e5e2b308c10a3812d
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio4/vocals.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e986bb445f169d4391a56fcebc5f55e8c39bcff752f61ccbf477161b303a1f69
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio5/bass.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b99f55edf4947722b0e0ed6b76099827f21eae020edac341d20f798758b615c5
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio5/drums.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eefadf55c12881632f35cf171e522117cc7e19c18683ece2887d2f9d5f79cb3a
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio5/other.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1da07f1d5c02b66d7d9091cce8bd8e5ddbc3aea1bd1f3bcb1eea8ff36abbd60
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio5/vocals.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b1411e7c0e128bc457873e011de7896f191f0bdee89c789e73a3d8c67df75f6
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio6/bass.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8dcafd9ecc1a257897feacd719ce27359c48fd0115d6593d1b747ec2a32613c1
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio6/drums.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cf62a180ea5d3939eec617a6dfa71ec8cc4a8cfe9fe90d282faf2aa2b698113
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio6/other.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bca9d5f51d1118a5c48228fd13825479e386bd3eedb929cdc3b5f71bb5db5f25
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/audios/htdemucs/audio6/vocals.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cc00cd8b79e898899f6d24a4216c7ccae407760bd2750746f652e07df8c2233
+size 5292044

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/metadata/audio1.json

@@ -0,0 +1 @@
+{"prompt": "a beautiful song"}

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/metadata/audio2.json

@@ -0,0 +1 @@
+{"prompt": "a beautiful song"}

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/metadata/audio3.json

@@ -0,0 +1 @@
+{"prompt": "a beautiful song"}

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/metadata/audio4.json

@@ -0,0 +1 @@
+{"prompt": "a beautiful song"}

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/metadata/audio5.json

@@ -0,0 +1 @@
+{"prompt": "a beautiful song"}

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/data/metadata/audio6.json

@@ -0,0 +1 @@
+{"prompt": "a beautiful song"}

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/dataset/dataloader.py

@@ -0,0 +1,191 @@
+import os
+import time
+import math
+import random
+import json
+import torch
+import torchaudio
+from torch.utils.data import Dataset, DataLoader, random_split
+
+from encodec import EncodecModel
+from encodec.utils import convert_audio
+
+
+class MusicDataset(Dataset):
+    def __init__(self, dataset_dir, sr, channels, min_duration, max_duration,
+                 sample_duration, aug_shift, device, composer=False):
+        super().__init__()
+        self.dataset_dir = dataset_dir
+        self.sr = sr
+        self.channels = channels
+        self.min_duration = min_duration 
+        self.max_duration = max_duration
+        self.sample_duration = sample_duration
+        self.aug_shift = aug_shift
+        self.device = device
+        self.model = EncodecModel.encodec_model_48khz().to(device=self.device)
+        self.audio_files_dir = f'{dataset_dir}/audios'
+        self.metadatas_dir = f'{dataset_dir}/metadata'
+        self.composer = composer
+        if composer:
+            self.demix_dir = f'{self.audio_files_dir}/htdemucs'
+        self.init_dataset()
+    
+    def get_duration_sec(self, file): 
+        wav, sr = torchaudio.load(file)
+        duration_sec = wav.shape[-1] / sr
+        return duration_sec
+    
+    def filter(self, audio_files, durations):
+        keep = []
+        self.audio_files = []
+        for i in range(len(audio_files)):
+            filepath = audio_files[i]
+            if durations[i] / self.sr < self.min_duration:
+                continue
+            if durations[i] / self.sr >= self.max_duration:
+                continue
+            keep.append(i)
+            self.audio_files.append(filepath)
+        self.durations = [durations[i] for i in keep] # in (s)
+        duration_tensor = torch.tensor(self.durations)
+        self.cumsum = torch.cumsum(duration_tensor, dim=0) # in (s)
+    
+    def init_dataset(self):
+        audio_files = os.listdir(self.audio_files_dir)
+        audio_files = [f'{self.audio_files_dir}/{file}' for file in audio_files if file.endswith('.wav') or file.endswith('.mp3')]
+        durations = [self.get_duration_sec(file) for file in audio_files]
+        self.filter(audio_files=audio_files, durations=durations)
+    
+    def get_index_offset(self, item):
+        half_interval = self.sample_duration // 2
+        shift = random.randint(-half_interval, half_interval) if self.aug_shift else 0
+        offset = item * self.sample_duration + shift
+        midpoint = offset + half_interval
+        assert 0 <= midpoint < self.cumsum[-1], f'Midpoint {midpoint} of item beyond total length {self.cumsum[-1]}'
+        index = torch.searchsorted(self.cumsum, midpoint)
+        start, end = self.cumsum[index-1] if index > 0 else 0.0, self.cumsum[index]
+        assert start <= midpoint <= end, f'Midpoint {midpoint} not inside interval [{start}, {end}] for index {index}'
+        if offset > end - self.sample_duration:
+            offset = max(start, offset - half_interval)
+        elif offset < start:
+            offset = min(end - self.sample_duration, offset + half_interval)
+        assert start <= offset <= end - self.sample_duration, f'Offset {offset} not in [{start}, {end} for index {index}]'
+        offset = offset - start
+        return index, offset
+    
+    def get_song_chunk(self, index, offset):
+        audio_file_path = self.audio_files[index]
+        song_name = os.path.splitext(os.path.basename(audio_file_path))[0]
+        wav, sr = torchaudio.load(audio_file_path)
+        
+        start_sample = int(offset * sr)
+        end_sample = start_sample + int(self.sample_duration * sr)
+        chunk = wav[:, start_sample:end_sample]
+        if self.composer:
+            demix_chunks = {}
+            demix_file_dict = {'bass': f'{self.demix_dir}/{song_name}/bass.wav',
+                               'drums': f'{self.demix_dir}/{song_name}/drums.wav',
+                               'other': f'{self.demix_dir}/{song_name}/other.wav'}
+            for key, value in demix_file_dict.items():
+                demix_chunk, demix_sr = torchaudio.load(value)
+                start_sample = int(offset * demix_sr)
+                end_sample = start_sample + int(self.sample_duration * demix_sr)
+                demix_chunks[key] = {'demix_chunk': demix_chunk[:, start_sample:end_sample],
+                                     'demix_sr': demix_sr}
+            
+            return chunk, sr, demix_chunks
+        #chunk = chunk.unsqueeze(0)
+        
+        return chunk, sr
+    
+    def __len__(self):
+        return len(self.durations)
+        
+    def __getitem__(self, item):
+        index, offset = self.get_index_offset(item)
+        if self.composer:
+            chunk, sr, demix_chunks = self.get_song_chunk(item, offset)
+        else:
+            chunk, sr = self.get_song_chunk(item, offset)
+        song_name = os.path.splitext(os.path.basename(self.audio_files[index]))[0]
+        if os.path.exists(f'{self.metadatas_dir}/{song_name}.json'):
+            with open(f'{self.metadatas_dir}/{song_name}.json', 'r') as file:
+                metadata = json.load(file)
+        chunk = convert_audio(chunk, sr, self.model.sample_rate, self.model.channels)
+        chunk = chunk.unsqueeze(0).to(device=self.device)
+        with torch.no_grad():
+            encoded_frames = self.model.encode(chunk)
+        chunk = chunk.mean(0, keepdim=True)
+        codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1)
+        codes = codes.transpose(0, 1)
+        emb = self.model.quantizer.decode(codes)
+        emb = emb.to(self.device)
+        
+        demix_embs = None
+        if self.composer:
+            demix_embs = {}
+            for key, value in demix_chunks.items():
+                demix_chunk = value['demix_chunk']
+                demix_sr = value['demix_sr']
+                demix_chunk = convert_audio(demix_chunk, demix_sr, self.model.sample_rate, self.model.channels)
+                demix_chunk = demix_chunk.unsqueeze(0).to(device=self.device)
+                with torch.no_grad():
+                    demix_encoded_frames = self.model.encode(demix_chunk)
+                demix_chunk = demix_chunk.mean(0, keepdim=True)
+                demix_codes = torch.cat([encoded[0] for encoded in demix_encoded_frames], dim=-1)
+                demix_codes = demix_codes.transpose(0, 1)
+                demix_emb = self.model.quantizer.decode(demix_codes)
+                demix_emb = demix_emb.to(self.device)
+                demix_embs[key] = demix_emb
+        
+        return chunk, metadata, emb, demix_embs
+
+def collate(batch):
+    device = 'cuda' if torch.cuda.is_available else 'cpu'
+    audio, data, emb, demix_embs = zip(*batch)
+    audio = torch.cat(audio, dim=0)
+    emb = torch.cat(emb, dim=0)
+    demix_embs_dict = None
+    metadata = [d for d in data]
+    keys = demix_embs[0].keys()
+    tensors_dict = {key: [] for key in keys}
+    if demix_embs is not None:
+        for d in demix_embs:
+            for key in keys:
+                tensors_dict[key].append(d[key])
+        demix_embs_dict = {key: torch.cat(tensors_dict[key], dim=0) for key in keys}
+        return (emb, metadata, demix_embs_dict)
+
+    return (emb, metadata, demix_embs_dict)
+
+def get_dataloader(dataset_folder, batch_size: int = 50, shuffle: bool = True):
+    dataset = MusicDataset(dataset_folder)
+
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate)
+    return dataloader
+
+
+def get_dataloaders(dataset_dir, sr, channels, min_duration, max_duration, sample_duration, 
+                    aug_shift, batch_size: int = 50, shuffle: bool = True, split_ratio=0.8, device='cpu',
+                    composer=False):
+    if not isinstance(dataset_dir, tuple):
+        dataset = MusicDataset(dataset_dir=dataset_dir, sr=sr, channels=channels,
+                               min_duration=min_duration, max_duration=max_duration, sample_duration=sample_duration,
+                               aug_shift=aug_shift, device=device, composer=composer)
+        # Split the dataset into train and validation
+        train_size = int(split_ratio * len(dataset))
+        val_size = len(dataset) - train_size
+        train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
+    else:
+        train_dir, valid_dir = dataset_dir
+        train_dataset = MusicDataset(dataset_dir=train_dir, sr=sr, channels=channels,
+                                     min_duration=min_duration, max_duration=max_duration, sample_duration=sample_duration,
+                                     aug_shift=aug_shift, device=device, composer=composer)
+        val_dataset = MusicDataset(dataset_dir=valid_dir, sr=sr, channels=channels,
+                                     min_duration=min_duration, max_duration=max_duration, sample_duration=sample_duration,
+                                     aug_shift=aug_shift, device=device, composer=composer)    
+    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate, drop_last=True)
+    val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate, drop_last=True)
+
+    return train_dataloader, val_dataloader

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/dataset/demix.py

@@ -0,0 +1,51 @@
+import sys
+import subprocess
+import torch
+
+from pathlib import Path
+from typing import List, Union
+
+
+def demix(paths: List[Path], demix_dir: Path, device: Union[str, torch.device]):
+  """Demixes the audio file into its sources."""
+  todos = []
+  demix_paths = []
+  for path in paths:
+    out_dir = demix_dir / 'htdemucs' / path.stem
+    demix_paths.append(out_dir)
+    if out_dir.is_dir():
+      if (
+        (out_dir / 'bass.wav').is_file() and
+        (out_dir / 'drums.wav').is_file() and
+        (out_dir / 'other.wav').is_file() and
+        (out_dir / 'vocals.wav').is_file()
+      ):
+        continue
+    todos.append(path)
+
+  existing = len(paths) - len(todos)
+  print(f'=> Found {existing} tracks already demixed, {len(todos)} to demix.')
+
+  if todos:
+    subprocess.run(
+      [
+        sys.executable, '-m', 'demucs.separate',
+        '--out', demix_dir.as_posix(),
+        '--name', 'htdemucs',
+        '--device', str(device),
+        *[path.as_posix() for path in todos],
+      ],
+      check=True,
+    )
+
+  return demix_paths
+
+def find_audio_files(directory: Path):
+    return list(directory.rglob('*.mp3'))  
+
+if __name__ == '__main__':
+    dataset_dir = '/home/keito/data/audios'
+    audio_dir = Path(dataset_dir)
+    audio_paths = find_audio_files(audio_dir)
+    demixed_paths = demix(audio_paths, audio_dir, "cuda") 
+    print('demixed_paths', demixed_paths)

JEN-1-COMPOSER-pytorch-main/JEN-1-COMPOSER-pytorch-main/generation.py

@@ -0,0 +1,212 @@
+import torch
+import random
+import numpy as np
+import math
+
+from utils.script_util import create_multi_conditioner, load_checkpoint
+from utils.config import Config
+
+from encodec import EncodecModel
+from encodec.utils import convert_audio
+
+from jen1.diffusion.gdm.gdm import GaussianDiffusion
+from jen1.model.model import UNetCFG1d
+from jen1.diffusion.gdm.noise_schedule import get_beta_schedule
+
+class Jen1():
+    def __init__(self, 
+                 ckpt_path, 
+                 device='cuda' if torch.cuda.is_available() else 'cpu',
+                 sample_rate = 48000, 
+                 cross_attn_cond_ids=['prompt'], 
+                 global_cond_ids= [],
+                 input_concat_ids= ['masked_input', 'mask']):
+        self.ckpt_path = ckpt_path
+        self.device = device
+        self.sample_rate = sample_rate
+        self.config = Config
+        self.conditioner = create_multi_conditioner(self.config.conditioner_config)     
+        self.cross_attn_cond_ids = cross_attn_cond_ids
+        self.global_cond_ids = global_cond_ids
+        self.input_concat_ids = input_concat_ids 
+        
+        self.audio_encoder = EncodecModel.encodec_model_48khz()
+        
+    def get_model_and_diffusion(self, steps, use_gdm):
+        if use_gdm:
+            diffusion_config = self.config.diffusion_config.gaussian_diffusion
+        else: 
+            diffusion_config = self.config.diffusion_config.variational_diffusion
+        model_config = self.config.model_config
+        
+        if use_gdm:
+            betas, alphas = get_beta_schedule(diffusion_config.noise_schedule, diffusion_config.steps)
+            betas = betas.to(self.device)
+            betas = betas.to(torch.float32)
+            if alphas is not None:
+                alphas.to(self.device)
+                alphas = alphas.to(torch.float32)
+            diffusion = GaussianDiffusion(steps=diffusion_config.steps, betas=betas, alphas=alphas,
+                                          objective=diffusion_config.objective, loss_type=diffusion_config.loss_type,
+                                          device=self.device, cfg_dropout_proba=diffusion_config.cfg_dropout_proba,
+                                          embedding_scale=diffusion_config.embedding_scale,
+                                          batch_cfg=diffusion_config.batch_cfg, scale_cfg=diffusion_config.scale_cfg,
+                                          sampling_timesteps=steps, use_fp16=False)
+            
+        config_dict = {k: v for k, v in model_config.__dict__.items() if not k.startswith('__') and not callable(v)}
+        context_embedding_features = config_dict.pop('context_embedding_features', None)
+        context_embedding_max_length = config_dict.pop('context_embedding_max_length', None)
+        
+        model = UNetCFG1d(context_embedding_features=context_embedding_features, 
+                     context_embedding_max_length=context_embedding_max_length, 
+                     **config_dict).to(self.device)
+        
+        #model, _, _, _ = load_checkpoint(self.ckpt_path, model)
+        model.eval()
+        diffusion.eval()
+        
+        return diffusion, model
+        
+    def generate(self, prompt, seed=-1, steps=100, batch_size=1, seconds=30, use_gdm=True,
+                 task='text_guided', init_audio=None, init_audio_sr=None, inpainting_scope=None):
+        
+        seed = seed if seed != -1 else np.random.randint(0, 2**32 -1)
+        torch.manual_seed(seed)
+        self.batch_size = batch_size
+        
+        diffusion, model = self.get_model_and_diffusion(steps, use_gdm)
+                
+        if init_audio is not None and init_audio.size() != 3:
+            init_audio = init_audio.repeat(batch_size, 1, 1)
+        
+        if init_audio is None:
+            flag = True
+            sample_length = seconds * self.sample_rate
+            shape = (batch_size, self.audio_encoder.channels, sample_length)
+            init_audio = torch.zeros(shape)
+            init_audio_sr = self.sample_rate
+        
+        init_audio = convert_audio(init_audio, init_audio_sr, self.sample_rate, self.audio_encoder.channels)
+
+        if task == 'text_guided':
+            mask = self.get_mask(sample_length, 0, seconds, batch_size)
+            causal = False
+        elif task == 'music_inpaint':
+            mask = self.get_mask(sample_length, inpainting_scope[0], inpainting_scope[1], batch_size)
+            causal = False
+        elif task == 'music_cont':
+            cont_length = sample_length - init_audio.size(2)
+            cont_start = init_audio.size(2)
+            mask = self.get_mask(sample_length, cont_start/self.sample_rate, seconds, batch_size)
+            cont_audio = torch.randn(batch_size, self.audio_encoder.channels, cont_length, device=self.device)
+            cont_audio = cont_audio * mask[:, cont_start:]   
+            init_audio = torch.cat([init_audio, cont_audio], dim=2) 
+            causal = True   
+        
+        with torch.no_grad():
+            init_emb = self.get_emb(init_audio).to(self.device)
+            emb_shape = init_emb.shape
+            mask = mask.to(self.device)
+            
+            mask = torch.nn.functional.interpolate(mask, size=(emb_shape[2]))
+            masked_emb = init_emb * mask
+            if flag:
+                init_emb = None
+            batch_metadata = [{'prompt': prompt} for _ in range(batch_size)]
+            conditioning = self.conditioner(batch_metadata, self.device)
+            conditioning['masked_input'] = masked_emb
+            conditioning['mask'] = mask
+            conditioning = self.get_conditioning(conditioning)
+            
+            sample_embs = diffusion.sample(model, emb_shape, conditioning, causal, init_data=init_emb)
+            sample_embs = sample_embs.to('cpu')
+            samples = self.audio_encoder.decoder(sample_embs)
+        
+        return samples
+    
+    def get_mask(self, sample_size, start, end, batch_size):
+        masks = []
+        maskstart = math.floor(start * self.sample_rate)
+        maskend = math.ceil(end * self.sample_rate)
+        mask = torch.ones((1, 1, sample_size))
+        mask[:, :, maskstart:maskend] = 0
+        masks.append(mask)
+        mask = torch.concat(masks * batch_size, dim=0)
+    
+        return mask
+    
+    def get_emb(self, audio):
+        encoded_frames = self.audio_encoder.encode(audio)
+        codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1)
+        codes = codes.transpose(0, 1)
+        emb = self.audio_encoder.quantizer.decode(codes)
+        return emb    
+            
+    def get_conditioning(self, cond):
+        cross_attention_input = None
+        cross_attention_masks = None
+        global_cond = None
+        input_concat_cond = None
+
+        if len(self.cross_attn_cond_ids) > 0:
+            # Concatenate all cross-attention inputs over the sequence dimension
+            # Assumes that the cross-attention inputs are of shape (batch, seq, channels)
+            cross_attention_input = torch.cat([cond[key][0] for key in self.cross_attn_cond_ids], dim=1)
+            cross_attention_masks = torch.cat([cond[key][1] for key in self.cross_attn_cond_ids], dim=1)
+
+        if len(self.global_cond_ids) > 0:
+            # Concatenate all global conditioning inputs over the channel dimension
+            # Assumes that the global conditioning inputs are of shape (batch, channels)
+            global_cond = torch.cat([cond[key][0] for key in self.global_cond_ids], dim=-1)
+            if len(global_cond.shape) == 3:
+                global_cond = global_cond.squeeze(1)
+        
+        if len(self.input_concat_ids) > 0:
+            concated_tensors = []
+            for key in self.input_concat_ids:
+                tensor = cond[key][0]
+                
+                if tensor.ndim == 2:
+                    tensor = tensor.unsqueeze(0)
+                    tensor = tensor.expand(self.batch_size, -1, -1)
+                
+                concated_tensors.append(tensor)
+            # Concatenate all input concat conditioning inputs over the channel dimension
+            # Assumes that the input concat conditioning inputs are of shape (batch, channels, seq)
+            #input_concat_cond = torch.cat([cond[key][0] for key in self.input_concat_ids], dim=1)
+            #For some reason, the BATCH component is removed. I don't know why.
+            input_concat_cond = torch.cat(concated_tensors, dim=1)
+            
+        return {
+            "cross_attn_cond": cross_attention_input,
+            "cross_attn_masks": cross_attention_masks,
+            "global_cond": global_cond,
+            "input_concat_cond": input_concat_cond
+        }
+        
+def save_audio_tensor(audio_tensor: torch.Tensor, file_path: str, sample_rate: int = 48000):
+    print(f'Saving audio to {file_path}')
+    """
+    Saves an audio tensor to a file.
+    Params:
+        audio_tensor: torch.Tensor, The audio data to save.
+        file_path: str, The path to the file where the audio will be saved.
+        sample_rate: int, The sample rate of the audio data.
+    Returns:
+        None
+    """
+    # Ensure the tensor is on the CPU before saving
+    audio_tensor = audio_tensor.detach()
+    print(f'audio_tensor.shape: {audio_tensor.shape}')
+    if audio_tensor.ndim == 3:
+        audio_tensor = audio_tensor.squeeze(0)  # Remove the batch dimension
+    # Use torchaudio to save the tensor as an audio file
+    import torchaudio
+    torchaudio.save(file_path, audio_tensor, sample_rate)
+    print(f'Saved audio to {file_path}')
+        
+if __name__ == '__main__':
+    jen1 = Jen1(ckpt_path=None)
+    prompt = 'a beautiful song'
+    samples = jen1.generate(prompt=prompt)
+    save_audio_tensor(samples, 'samples.wav')