v1.2

app.py

@@ -27,22 +27,16 @@ from diffrhythm.infer.infer import inference
 
 MAX_SEED = np.iinfo(np.int32).max
 device='cuda'
-cfm, cfm_full, tokenizer, muq, vae = prepare_model(device)
+cfm, tokenizer, muq, vae, eval_model, eval_muq = prepare_model(device)
 cfm = torch.compile(cfm)
-cfm_full = torch.compile(cfm_full)
 
-@spaces.GPU(duration=40)
-def infer_music(lrc, ref_audio_path, text_prompt, current_prompt_type, seed=42, randomize_seed=False, steps=32, cfg_strength=4.0, file_type='wav', odeint_method='euler', Music_Duration='95s', device='cuda'):
-    if Music_Duration == '95s':
-        max_frames = 2048
-        cfm_model = cfm
-    else:
-        max_frames = 6144
-        cfm_model = cfm_full
+@spaces.GPU
+def infer_music(lrc, ref_audio_path, text_prompt, current_prompt_type, seed=42, randomize_seed=False, steps=32, cfg_strength=4.0, file_type='wav', odeint_method='euler', preference_infer="quality first", edit=False, edit_segments=None, device='cuda'):
+    max_frames = 2048
+    sway_sampling_coef = -1 if steps < 32 else None
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     torch.manual_seed(seed)
-    sway_sampling_coef = -1 if steps < 32 else None
     vocal_flag = False
     try:
         lrc_prompt, start_time = get_lrc_token(max_frames, lrc, tokenizer, device)
@@ -53,9 +47,16 @@ def infer_music(lrc, ref_audio_path, text_prompt, current_prompt_type, seed=42,
     except Exception as e:
         raise gr.Error(f"Error: {str(e)}")
     negative_style_prompt = get_negative_style_prompt(device)
-    latent_prompt = get_reference_latent(device, max_frames)
-    generated_song = inference(cfm_model=cfm_model, 
+    latent_prompt, pred_frames = get_reference_latent(device, max_frames, edit, edit_segments, ref_audio_path, vae)
+    
+    if preference_infer == "quality first":
+        batch_infer_num = 5
+    else:
+        batch_infer_num = 1
+    generated_song = inference(cfm_model=cfm, 
                                vae_model=vae, 
+                               eval_model=eval_model,
+                               eval_muq=eval_muq,
                                cond=latent_prompt, 
                                text=lrc_prompt, 
                                duration=max_frames, 
@@ -68,6 +69,8 @@ def infer_music(lrc, ref_audio_path, text_prompt, current_prompt_type, seed=42,
                                file_type=file_type,
                                vocal_flag=vocal_flag,
                                odeint_method=odeint_method,
+                               pred_frames=pred_frames,
+                               batch_infer_num=batch_infer_num,
                                )
     return generated_song
 
@@ -234,8 +237,8 @@ with gr.Blocks(css=css) as demo:
     - If loading audio result is slow, you can select Output Format as mp3 in Advanced Settings.                                
     
                         """)
-                    Music_Duration = gr.Radio(["95s", "285s"], label="Music Duration", value="95s")
-                    
+                    # Music_Duration = gr.Radio(["95s", "285s"], label="Music Duration", value="95s")
+                    preference_infer = gr.Radio(["quality first", "speed first"], label="Preference", value="quality first")
                     lyrics_btn = gr.Button("Generate", variant="primary")
                     audio_output = gr.Audio(label="Audio Result", type="filepath", elem_id="audio_output")
                     with gr.Accordion("Advanced Settings", open=False):
@@ -266,6 +269,12 @@ with gr.Blocks(css=css) as demo:
                                     interactive=True,
                                     elem_id="step_slider"
                                 )
+                        edit = gr.Checkbox(label="edit", value=False)
+                        edit_segeditments = gr.Textbox(
+                            label="Edit Segments",
+                            placeholder="Time segments to edit (in seconds). Format: `[[start1,end1],...]",
+                            )
+                        
                         odeint_method = gr.Radio(["euler", "midpoint", "rk4","implicit_adams"], label="ODE Solver", value="euler")                        
                         file_type = gr.Dropdown(["wav", "mp3", "ogg"], label="Output Format", value="wav")
 
@@ -409,7 +418,7 @@ with gr.Blocks(css=css) as demo:
     
     lyrics_btn.click(
         fn=infer_music,
-        inputs=[lrc, audio_prompt, text_prompt, current_prompt_type, seed, randomize_seed, steps, cfg_strength, file_type, odeint_method, Music_Duration],
+        inputs=[lrc, audio_prompt, text_prompt, current_prompt_type, seed, randomize_seed, steps, cfg_strength, file_type, odeint_method, preference_infer, edit, edit_segments],
         outputs=audio_output
     )
 

diffrhythm/config/{diffrhythm-1b.json → config.json}

@@ -2,7 +2,7 @@
     "model_type": "diffrhythm",
     "model": {
         "dim": 2048,
-        "depth": 16,
+        "depth": 16, 
         "heads": 32,
         "ff_mult": 4,
         "text_dim": 512,

diffrhythm/infer/infer.py

@@ -2,82 +2,51 @@ import torch
 import torchaudio
 from einops import rearrange
 import argparse
-import json
 import os
-from tqdm import tqdm
+import time
 import random
+
+import torch
+import torchaudio
 import numpy as np
-import time
+from einops import rearrange
 import io
 import pydub
 
 from diffrhythm.infer.infer_utils import (
-    get_reference_latent,
+    decode_audio,
     get_lrc_token,
-    get_audio_style_prompt,
+    get_negative_style_prompt,
+    get_reference_latent,
+    get_style_prompt,
     prepare_model,
-    get_negative_style_prompt
+    eval_song,
 )
 
-def decode_audio(latents, vae_model, chunked=False, overlap=32, chunk_size=128):
-    downsampling_ratio = 2048
-    io_channels = 2
-    if not chunked:
-        # default behavior. Decode the entire latent in parallel
-        return vae_model.decode_export(latents)
-    else:
-        # chunked decoding
-        hop_size = chunk_size - overlap
-        total_size = latents.shape[2]
-        batch_size = latents.shape[0]
-        chunks = []
-        i = 0
-        for i in range(0, total_size - chunk_size + 1, hop_size):
-            chunk = latents[:,:,i:i+chunk_size]
-            chunks.append(chunk)
-        if i+chunk_size != total_size:
-            # Final chunk
-            chunk = latents[:,:,-chunk_size:]
-            chunks.append(chunk)
-        chunks = torch.stack(chunks)
-        num_chunks = chunks.shape[0]
-        # samples_per_latent is just the downsampling ratio
-        samples_per_latent = downsampling_ratio
-        # Create an empty waveform, we will populate it with chunks as decode them
-        y_size = total_size * samples_per_latent
-        y_final = torch.zeros((batch_size,io_channels,y_size)).to(latents.device)
-        for i in range(num_chunks):
-            x_chunk = chunks[i,:]
-            # decode the chunk
-            y_chunk = vae_model.decode_export(x_chunk)
-            # figure out where to put the audio along the time domain
-            if i == num_chunks-1:
-                # final chunk always goes at the end
-                t_end = y_size
-                t_start = t_end - y_chunk.shape[2]
-            else:
-                t_start = i * hop_size * samples_per_latent
-                t_end = t_start + chunk_size * samples_per_latent
-            #  remove the edges of the overlaps
-            ol = (overlap//2) * samples_per_latent
-            chunk_start = 0
-            chunk_end = y_chunk.shape[2]
-            if i > 0:
-                # no overlap for the start of the first chunk
-                t_start += ol
-                chunk_start += ol
-            if i < num_chunks-1:
-                # no overlap for the end of the last chunk
-                t_end -= ol
-                chunk_end -= ol
-            # paste the chunked audio into our y_final output audio
-            y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]
-        return y_final
-
-def inference(cfm_model, vae_model, cond, text, duration, style_prompt, negative_style_prompt, steps, cfg_strength, sway_sampling_coef, start_time, file_type, vocal_flag, odeint_method):
 
+def inference(
+    cfm_model,
+    vae_model,
+    eval_model,
+    eval_muq,
+    cond,
+    text,
+    duration,
+    style_prompt,
+    negative_style_prompt,
+    steps,
+    cfg_strength,
+    sway_sampling_coef,
+    start_time,
+    file_type,
+    vocal_flag,
+    odeint_method,
+    pred_frames,
+    batch_infer_num,
+    chunked=True,
+):
     with torch.inference_mode():
-        generated, _ = cfm_model.sample(
+        latents, _ = cfm_model.sample(
             cond=cond,
             text=text,
             duration=duration,
@@ -89,17 +58,27 @@ def inference(cfm_model, vae_model, cond, text, duration, style_prompt, negative
             start_time=start_time,
             vocal_flag=vocal_flag,
             odeint_method=odeint_method,
+            latent_pred_segments=pred_frames,
+            batch_infer_num=batch_infer_num
         )
-        
-        generated = generated.to(torch.float32)
-        latent = generated.transpose(1, 2) # [b d t]
-        output = decode_audio(latent, vae_model, chunked=False)
 
-        # Rearrange audio batch to a single sequence
-        output = rearrange(output, "b d n -> d (b n)")
-        output_tensor = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).cpu()
+        outputs = []
+        for latent in latents:
+            latent = latent.to(torch.float32)
+            latent = latent.transpose(1, 2)  # [b d t]
+
+            output = decode_audio(latent, vae_model, chunked=chunked)
+
+            # Rearrange audio batch to a single sequence
+            output = rearrange(output, "b d n -> d (b n)")
+            
+            outputs.append(output)
+        if batch_infer_num > 1:
+            generated_song = eval_song(eval_model, eval_muq, outputs)
+        else:
+            generated_song = outputs[0]
+        output_tensor = generated_song.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).cpu()
         output_np = output_tensor.numpy().T.astype(np.float32)
-        
         if file_type == 'wav':
             return (44100, output_np)
         else:
@@ -111,52 +90,140 @@ def inference(cfm_model, vae_model, cond, text, duration, style_prompt, negative
             else:
                 song.export(buffer, format="ogg", bitrate="320k")
             return buffer.getvalue()
-    
         
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--lrc-path', type=str, default="example/eg.lrc") # lyrics of target song
-    parser.add_argument('--ref-audio-path', type=str, default="example/eg.mp3") # reference audio as style prompt for target song
-    parser.add_argument('--audio-length', type=int, default=95) # length of target song
-    parser.add_argument('--output-dir', type=str, default="example/output")
+    parser.add_argument(
+        "--lrc-path",
+        type=str,
+        help="lyrics of target song",
+    )  # lyrics of target song
+    parser.add_argument(
+        "--ref-prompt",
+        type=str,
+        help="reference prompt as style prompt for target song",
+        required=False,
+    )  # reference prompt as style prompt for target song
+    parser.add_argument(
+        "--ref-audio-path",
+        type=str,
+        help="reference audio as style prompt for target song",
+        required=False,
+    )  # reference audio as style prompt for target song
+    parser.add_argument(
+        "--chunked",
+        action="store_true",
+        help="whether to use chunked decoding",
+    )  # whether to use chunked decoding
+    parser.add_argument(
+        "--audio-length",
+        type=int,
+        default=95,
+        choices=[95, 285],
+        help="length of generated song",
+    )  # length of target song
+    parser.add_argument(
+        "--repo-id", type=str, default="ASLP-lab/DiffRhythm-base", help="target model"
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="infer/example/output",
+        help="output directory fo generated song",
+    )  # output directory of target song
+    parser.add_argument(
+        "--edit",
+        action="store_true",
+        help="whether to open edit mode",
+    )  # edit flag
+    parser.add_argument(
+        "--ref-song",
+        type=str,
+        required=False,
+        help="reference prompt as latent prompt for editing",
+    )  # reference prompt as latent prompt for editing
+    parser.add_argument(
+        "--edit-segments",
+        type=str,
+        required=False,
+        help="edit segments o target song",
+    )  # edit segments o target song
     args = parser.parse_args()
-    
-    device = 'cuda'
-    
+
+    assert (
+        args.ref_prompt or args.ref_audio_path
+    ), "either ref_prompt or ref_audio_path should be provided"
+    assert not (
+        args.ref_prompt and args.ref_audio_path
+    ), "only one of them should be provided"
+    if args.edit:
+        assert (
+            args.ref_song and args.edit_segments
+        ), "reference song and edit segments should be provided for editing"
+
+    device = "cpu"
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.mps.is_available():
+        device = "mps"
+
     audio_length = args.audio_length
     if audio_length == 95:
         max_frames = 2048
     elif audio_length == 285:
         max_frames = 6144
+
+    cfm, tokenizer, muq, vae, eval_model, eval_muq = prepare_model(max_frames, device, repo_id=args.repo_id)
+
+    if args.lrc_path:
+        with open(args.lrc_path, "r", encoding='utf-8') as f:
+            lrc = f.read()
+    else:
+        lrc = ""
+    lrc_prompt, start_time = get_lrc_token(max_frames, lrc, tokenizer, device)
+
+    if args.ref_audio_path:
+        style_prompt = get_style_prompt(muq, args.ref_audio_path)
+    else:
+        style_prompt = get_style_prompt(muq, prompt=args.ref_prompt)
+
+    negative_style_prompt = get_negative_style_prompt(device)
+
+    latent_prompt, pred_frames = get_reference_latent(device, max_frames, args.edit, args.edit_segments, args.ref_song, vae)
+
+    s_t = time.time()
+    generated_songs = inference(
+        cfm_model=cfm,
+        vae_model=vae,
+        cond=latent_prompt,
+        text=lrc_prompt,
+        duration=max_frames,
+        style_prompt=style_prompt,
+        negative_style_prompt=negative_style_prompt,
+        start_time=start_time,
+        pred_frames=pred_frames,
+        chunked=args.chunked,
+    )
     
-    cfm, tokenizer, muq, vae = prepare_model(device)
-    
-    with open(args.lrc_path, 'r') as f:
-        lrc = f.read()
-    lrc_prompt, start_time = get_lrc_token(lrc, tokenizer, device)
-    
-    style_prompt = get_audio_style_prompt(muq, args.ref_audio_path)
     
-    negative_style_prompt = get_negative_style_prompt(device)
     
-    latent_prompt = get_reference_latent(device, max_frames)
+    generated_song = eval_song(eval_model, eval_muq, generated_songs)
     
-    s_t = time.time()
-    generated_song = inference(cfm_model=cfm, 
-                               vae_model=vae, 
-                               cond=latent_prompt, 
-                               text=lrc_prompt, 
-                               duration=max_frames, 
-                               style_prompt=style_prompt,
-                               negative_style_prompt=negative_style_prompt,
-                               start_time=start_time
-                               )
+    # Peak normalize, clip, convert to int16, and save to file
+    generated_song = (
+        generated_song.to(torch.float32)
+        .div(torch.max(torch.abs(generated_song)))
+        .clamp(-1, 1)
+        .mul(32767)
+        .to(torch.int16)
+        .cpu()
+    )
     e_t = time.time() - s_t
-    print(f"inference cost {e_t} seconds")
-    
+    print(f"inference cost {e_t:.2f} seconds")
     output_dir = args.output_dir
     os.makedirs(output_dir, exist_ok=True)
-    
+
     output_path = os.path.join(output_dir, "output.wav")
     torchaudio.save(output_path, generated_song, sample_rate=44100)
-    

diffrhythm/infer/infer_utils.py

@@ -1,66 +1,308 @@
 import torch
 import librosa
+import torchaudio
 import random
 import json
-from muq import MuQMuLan
+from muq import MuQMuLan, MuQ
 from mutagen.mp3 import MP3
 import os
 import numpy as np
 from huggingface_hub import hf_hub_download
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+from safetensors.torch import load_file
+
 from diffrhythm.model import DiT, CFM
 
+def vae_sample(mean, scale):
+    stdev = torch.nn.functional.softplus(scale) + 1e-4
+    var = stdev * stdev
+    logvar = torch.log(var)
+    latents = torch.randn_like(mean) * stdev + mean
+
+    kl = (mean * mean + var - logvar - 1).sum(1).mean()
+
+    return latents, kl
+
+def normalize_audio(y, target_dbfs=0):
+    max_amplitude = torch.max(torch.abs(y))
+
+    target_amplitude = 10.0**(target_dbfs / 20.0)
+    scale_factor = target_amplitude / max_amplitude
+
+    normalized_audio = y * scale_factor
+
+    return normalized_audio
+
+def set_audio_channels(audio, target_channels):
+    if target_channels == 1:
+        # Convert to mono
+        audio = audio.mean(1, keepdim=True)
+    elif target_channels == 2:
+        # Convert to stereo
+        if audio.shape[1] == 1:
+            audio = audio.repeat(1, 2, 1)
+        elif audio.shape[1] > 2:
+            audio = audio[:, :2, :]
+    return audio
+
+class PadCrop(torch.nn.Module):
+    def __init__(self, n_samples, randomize=True):
+        super().__init__()
+        self.n_samples = n_samples
+        self.randomize = randomize
+
+    def __call__(self, signal):
+        n, s = signal.shape
+        start = 0 if (not self.randomize) else torch.randint(0, max(0, s - self.n_samples) + 1, []).item()
+        end = start + self.n_samples
+        output = signal.new_zeros([n, self.n_samples])
+        output[:, :min(s, self.n_samples)] = signal[:, start:end]
+        return output
+
+def prepare_audio(audio, in_sr, target_sr, target_length, target_channels, device):
+    
+    audio = audio.to(device)
+
+    if in_sr != target_sr:
+        resample_tf = T.Resample(in_sr, target_sr).to(device)
+        audio = resample_tf(audio)
+    if target_length is None:
+        target_length = audio.shape[-1]
+    audio = PadCrop(target_length, randomize=False)(audio)
+
+    # Add batch dimension
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0).unsqueeze(0)
+    elif audio.dim() == 2:
+        audio = audio.unsqueeze(0)
+
+    audio = set_audio_channels(audio, target_channels)
+
+    return audio
+
+def decode_audio(latents, vae_model, chunked=False, overlap=32, chunk_size=128):
+    downsampling_ratio = 2048
+    io_channels = 2
+    if not chunked:
+        return vae_model.decode_export(latents)
+    else:
+        # chunked decoding
+        hop_size = chunk_size - overlap
+        total_size = latents.shape[2]
+        batch_size = latents.shape[0]
+        chunks = []
+        i = 0
+        for i in range(0, total_size - chunk_size + 1, hop_size):
+            chunk = latents[:, :, i : i + chunk_size]
+            chunks.append(chunk)
+        if i + chunk_size != total_size:
+            # Final chunk
+            chunk = latents[:, :, -chunk_size:]
+            chunks.append(chunk)
+        chunks = torch.stack(chunks)
+        num_chunks = chunks.shape[0]
+        # samples_per_latent is just the downsampling ratio
+        samples_per_latent = downsampling_ratio
+        # Create an empty waveform, we will populate it with chunks as decode them
+        y_size = total_size * samples_per_latent
+        y_final = torch.zeros((batch_size, io_channels, y_size)).to(latents.device)
+        for i in range(num_chunks):
+            x_chunk = chunks[i, :]
+            # decode the chunk
+            y_chunk = vae_model.decode_export(x_chunk)
+            # figure out where to put the audio along the time domain
+            if i == num_chunks - 1:
+                # final chunk always goes at the end
+                t_end = y_size
+                t_start = t_end - y_chunk.shape[2]
+            else:
+                t_start = i * hop_size * samples_per_latent
+                t_end = t_start + chunk_size * samples_per_latent
+            #  remove the edges of the overlaps
+            ol = (overlap // 2) * samples_per_latent
+            chunk_start = 0
+            chunk_end = y_chunk.shape[2]
+            if i > 0:
+                # no overlap for the start of the first chunk
+                t_start += ol
+                chunk_start += ol
+            if i < num_chunks - 1:
+                # no overlap for the end of the last chunk
+                t_end -= ol
+                chunk_end -= ol
+            # paste the chunked audio into our y_final output audio
+            y_final[:, :, t_start:t_end] = y_chunk[:, :, chunk_start:chunk_end]
+        return y_final
+
+def encode_audio(audio, vae_model, chunked=False, overlap=32, chunk_size=128):
+    downsampling_ratio = 2048
+    latent_dim = 128
+    if not chunked:
+        # default behavior. Encode the entire audio in parallel
+        return vae_model.encode_export(audio)
+    else:
+        # CHUNKED ENCODING
+        # samples_per_latent is just the downsampling ratio (which is also the upsampling ratio)
+        samples_per_latent = downsampling_ratio
+        total_size = audio.shape[2] # in samples
+        batch_size = audio.shape[0]
+        chunk_size *= samples_per_latent # converting metric in latents to samples
+        overlap *= samples_per_latent # converting metric in latents to samples
+        hop_size = chunk_size - overlap
+        chunks = []
+        for i in range(0, total_size - chunk_size + 1, hop_size):
+            chunk = audio[:,:,i:i+chunk_size]
+            chunks.append(chunk)
+        if i+chunk_size != total_size:
+            # Final chunk
+            chunk = audio[:,:,-chunk_size:]
+            chunks.append(chunk)
+        chunks = torch.stack(chunks)
+        num_chunks = chunks.shape[0]
+        # Note: y_size might be a different value from the latent length used in diffusion training
+        # because we can encode audio of varying lengths
+        # However, the audio should've been padded to a multiple of samples_per_latent by now.
+        y_size = total_size // samples_per_latent
+        # Create an empty latent, we will populate it with chunks as we encode them
+        y_final = torch.zeros((batch_size,latent_dim,y_size)).to(audio.device)
+        for i in range(num_chunks):
+            x_chunk = chunks[i,:]
+            # encode the chunk
+            y_chunk = vae_model.encode_export(x_chunk)
+            # figure out where to put the audio along the time domain
+            if i == num_chunks-1:
+                # final chunk always goes at the end
+                t_end = y_size
+                t_start = t_end - y_chunk.shape[2]
+            else:
+                t_start = i * hop_size // samples_per_latent
+                t_end = t_start + chunk_size // samples_per_latent
+            #  remove the edges of the overlaps
+            ol = overlap//samples_per_latent//2
+            chunk_start = 0
+            chunk_end = y_chunk.shape[2]
+            if i > 0:
+                # no overlap for the start of the first chunk
+                t_start += ol
+                chunk_start += ol
+            if i < num_chunks-1:
+                # no overlap for the end of the last chunk
+                t_end -= ol
+                chunk_end -= ol
+            # paste the chunked audio into our y_final output audio
+            y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]
+        return y_final
 
 def prepare_model(device):
     # prepare cfm model
-    dit_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-base", filename="cfm_model.pt")
-    dit_full_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-full", filename="cfm_model.pt")
-    dit_config_path = "./diffrhythm/config/diffrhythm-1b.json"
+    
+    dit_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-1_2", filename="cfm_model.pt")
+    dit_config_path = "./diffrhythm/config/config.json"
     with open(dit_config_path) as f:
         model_config = json.load(f)
     dit_model_cls = DiT
     cfm = CFM(
-                transformer=dit_model_cls(**model_config["model"], use_style_prompt=True, max_pos=2048),
+                transformer=dit_model_cls(**model_config["model"], max_frames=2048),
                 num_channels=model_config["model"]['mel_dim'],
-                use_style_prompt=True
              )
     cfm = cfm.to(device)
     cfm = load_checkpoint(cfm, dit_ckpt_path, device=device, use_ema=False)
-    
-    cfm_full = CFM(
-                transformer=dit_model_cls(**model_config["model"], use_style_prompt=True, max_pos=6144),
-                num_channels=model_config["model"]['mel_dim'],
-                use_style_prompt=True
-             )
-    cfm_full = cfm_full.to(device)
-    cfm_full = load_checkpoint(cfm_full, dit_full_ckpt_path, device=device, use_ema=False)
-    
+
     # prepare tokenizer
     tokenizer = CNENTokenizer()
-    
+
     # prepare muq
-    muq = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large")
+    muq = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large", cache_dir="./pretrained")
     muq = muq.to(device).eval()
-    
+
     # prepare vae
     vae_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-vae", filename="vae_model.pt")
-    vae = torch.jit.load(vae_ckpt_path, map_location='cpu').to(device)
+    vae = torch.jit.load(vae_ckpt_path, map_location="cpu").to(device)
+    
     
-    return cfm, cfm_full, tokenizer, muq, vae
+    # prepare eval model
+    train_config = OmegaConf.load("./pretrained/eval.yaml")
+    checkpoint_path = "./pretrained/eval.safetensors"
+
+    eval_model = instantiate(train_config.generator).to(device).eval()
+    state_dict = load_file(checkpoint_path, device="cpu")
+    eval_model.load_state_dict(state_dict)
+
+    eval_muq = MuQ.from_pretrained("OpenMuQ/MuQ-large-msd-iter")
+    eval_muq = eval_muq.to(device).eval()
+
+    return cfm, tokenizer, muq, vae, eval_model, eval_muq
 
 
 # for song edit, will be added in the future
-def get_reference_latent(device, max_frames):
-    return torch.zeros(1, max_frames, 64).to(device)
+def get_reference_latent(device, max_frames, edit, pred_segments, ref_song, vae_model):
+    sampling_rate = 44100
+    downsample_rate = 2048
+    io_channels = 2
+    if edit:
+        input_audio, in_sr = torchaudio.load(ref_song)
+        input_audio = prepare_audio(input_audio, in_sr=in_sr, target_sr=sampling_rate, target_length=None, target_channels=io_channels, device=device)
+        input_audio = normalize_audio(input_audio, -6)
+        
+        with torch.no_grad():
+            latent = encode_audio(input_audio, vae_model, chunked=True) # [b d t]
+            mean, scale = latent.chunk(2, dim=1)
+            prompt, _ = vae_sample(mean, scale)
+            prompt = prompt.transpose(1, 2) # [b t d]
+        
+        pred_segments = json.loads(pred_segments)
+        # import pdb; pdb.set_trace()
+        pred_frames = []
+        for st, et in pred_segments:
+            sf = 0 if st == -1 else int(st * sampling_rate / downsample_rate)
+            # if st == -1:
+            #     sf = 0
+            # else:
+            #     sf = int(st * sampling_rate / downsample_rate )
+            
+            ef = max_frames if et == -1 else int(et * sampling_rate / downsample_rate)
+            # if et == -1:
+            #     ef = max_frames
+            # else:
+            #     ef = int(et * sampling_rate / downsample_rate )
+            pred_frames.append((sf, ef))
+        # import pdb; pdb.set_trace()
+        return prompt, pred_frames
+    else:
+        prompt = torch.zeros(1, max_frames, 64).to(device)
+        pred_frames = [(0, max_frames)]
+        return prompt, pred_frames
+
 
 def get_negative_style_prompt(device):
     file_path = "./src/negative_prompt.npy"
     vocal_stlye = np.load(file_path)
-    
-    vocal_stlye = torch.from_numpy(vocal_stlye).to(device) # [1, 512]
+
+    vocal_stlye = torch.from_numpy(vocal_stlye).to(device)  # [1, 512]
     vocal_stlye = vocal_stlye.half()
-    
+
     return vocal_stlye
 
+@torch.no_grad()
+def eval_song(eval_model, eval_muq, songs):
+    
+    resampled_songs = [torchaudio.functional.resample(song.mean(dim=0, keepdim=True), 44100, 24000) for song in songs]
+    ssl_list = []
+    for i in range(len(resampled_songs)):
+        output = eval_muq(resampled_songs[i], output_hidden_states=True)
+        muq_ssl = output["hidden_states"][6]
+        ssl_list.append(muq_ssl.squeeze(0))
+
+    ssl = torch.stack(ssl_list)
+    scores_g = eval_model(ssl)
+    score = torch.mean(scores_g, dim=1)
+    idx = score.argmax(dim=0)
+    
+    return songs[idx]
+    
+
+@torch.no_grad()
 def get_audio_style_prompt(model, wav_path):
     vocal_flag = False
     mulan = model
@@ -85,6 +327,8 @@ def get_audio_style_prompt(model, wav_path):
 
     return audio_emb, vocal_flag
 
+
+@torch.no_grad()
 def get_text_style_prompt(model, text_prompt):
     mulan = model
     
@@ -95,50 +339,88 @@ def get_text_style_prompt(model, text_prompt):
     return text_emb
 
 
+@torch.no_grad()
+def get_style_prompt(model, wav_path=None, prompt=None):
+    mulan = model
+
+    if prompt is not None:
+        return mulan(texts=prompt).half()
+
+    ext = os.path.splitext(wav_path)[-1].lower()
+    if ext == ".mp3":
+        meta = MP3(wav_path)
+        audio_len = meta.info.length
+    elif ext in [".wav", ".flac"]:
+        audio_len = librosa.get_duration(path=wav_path)
+    else:
+        raise ValueError("Unsupported file format: {}".format(ext))
+
+    if audio_len < 10:
+        print(
+            f"Warning: The audio file {wav_path} is too short ({audio_len:.2f} seconds). Expected at least 10 seconds."
+        )
+
+    assert audio_len >= 10
+
+    mid_time = audio_len // 2
+    start_time = mid_time - 5
+    wav, _ = librosa.load(wav_path, sr=24000, offset=start_time, duration=10)
+
+    wav = torch.tensor(wav).unsqueeze(0).to(model.device)
+
+    with torch.no_grad():
+        audio_emb = mulan(wavs=wav)  # [1, 512]
+
+    audio_emb = audio_emb
+    audio_emb = audio_emb.half()
+
+    return audio_emb
 
 def parse_lyrics(lyrics: str):
     lyrics_with_time = []
     lyrics = lyrics.strip()
-    for line in lyrics.split('\n'):
+    for line in lyrics.split("\n"):
         try:
             time, lyric = line[1:9], line[10:]
             lyric = lyric.strip()
-            mins, secs = time.split(':')
+            mins, secs = time.split(":")
             secs = int(mins) * 60 + float(secs)
             lyrics_with_time.append((secs, lyric))
         except:
             continue
     return lyrics_with_time
 
-class CNENTokenizer():
+
+class CNENTokenizer:
     def __init__(self):
-        with open('./diffrhythm/g2p/g2p/vocab.json', 'r') as file:
-            self.phone2id:dict = json.load(file)['vocab']
-        self.id2phone = {v:k for (k, v) in self.phone2id.items()}
+        with open("./diffrhythm/g2p/g2p/vocab.json", "r", encoding='utf-8') as file:
+            self.phone2id: dict = json.load(file)["vocab"]
+        self.id2phone = {v: k for (k, v) in self.phone2id.items()}
         from diffrhythm.g2p.g2p_generation import chn_eng_g2p
+
         self.tokenizer = chn_eng_g2p
+
     def encode(self, text):
         phone, token = self.tokenizer(text)
-        token = [x+1 for x in token]
+        token = [x + 1 for x in token]
         return token
+
     def decode(self, token):
-        return "|".join([self.id2phone[x-1] for x in token])
-    
+        return "|".join([self.id2phone[x - 1] for x in token])
+
+
 def get_lrc_token(max_frames, text, tokenizer, device):
 
     lyrics_shift = 0
     sampling_rate = 44100
     downsample_rate = 2048
     max_secs = max_frames / (sampling_rate / downsample_rate)
-   
-    pad_token_id = 0
+
     comma_token_id = 1
-    period_token_id = 2    
-    if text == "":
-        return torch.zeros((max_frames,), dtype=torch.long).unsqueeze(0).to(device), torch.tensor(0.).unsqueeze(0).to(device).half()
+    period_token_id = 2
 
     lrc_with_time = parse_lyrics(text)
-    
+
     modified_lrc_with_time = []
     for i in range(len(lrc_with_time)):
         time, line = lrc_with_time[i]
@@ -146,44 +428,49 @@ def get_lrc_token(max_frames, text, tokenizer, device):
         modified_lrc_with_time.append((time, line_token))
     lrc_with_time = modified_lrc_with_time
 
-    lrc_with_time = [(time_start, line) for (time_start, line) in lrc_with_time if time_start < max_secs]
-    # lrc_with_time = lrc_with_time[:-1] if len(lrc_with_time) >= 1 else lrc_with_time
-    
-    normalized_start_time = 0.
+    lrc_with_time = [
+        (time_start, line)
+        for (time_start, line) in lrc_with_time
+        if time_start < max_secs
+    ]
+    if max_frames == 2048:
+        lrc_with_time = lrc_with_time[:-1] if len(lrc_with_time) >= 1 else lrc_with_time
+
+    normalized_start_time = 0.0
 
     lrc = torch.zeros((max_frames,), dtype=torch.long)
 
     tokens_count = 0
     last_end_pos = 0
     for time_start, line in lrc_with_time:
-        tokens = [token if token != period_token_id else comma_token_id for token in line] + [period_token_id]
+        tokens = [
+            token if token != period_token_id else comma_token_id for token in line
+        ] + [period_token_id]
         tokens = torch.tensor(tokens, dtype=torch.long)
         num_tokens = tokens.shape[0]
 
         gt_frame_start = int(time_start * sampling_rate / downsample_rate)
-        
-        frame_shift = random.randint(int(lyrics_shift), int(lyrics_shift))
-        
+
+        frame_shift = random.randint(int(-lyrics_shift), int(lyrics_shift))
+
         frame_start = max(gt_frame_start - frame_shift, last_end_pos)
         frame_len = min(num_tokens, max_frames - frame_start)
 
-       
-
-        lrc[frame_start:frame_start + frame_len] = tokens[:frame_len]
+        lrc[frame_start : frame_start + frame_len] = tokens[:frame_len]
 
         tokens_count += num_tokens
-        last_end_pos = frame_start + frame_len   
-        
+        last_end_pos = frame_start + frame_len
+
     lrc_emb = lrc.unsqueeze(0).to(device)
-    
+
     normalized_start_time = torch.tensor(normalized_start_time).unsqueeze(0).to(device)
     normalized_start_time = normalized_start_time.half()
-    
+
     return lrc_emb, normalized_start_time
 
+
 def load_checkpoint(model, ckpt_path, device, use_ema=True):
-    if device == "cuda":
-        model = model.half()
+    model = model.half()
 
     ckpt_type = ckpt_path.split(".")[-1]
     if ckpt_type == "safetensors":
@@ -207,4 +494,4 @@ def load_checkpoint(model, ckpt_path, device, use_ema=True):
             checkpoint = {"model_state_dict": checkpoint}
         model.load_state_dict(checkpoint["model_state_dict"], strict=False)
 
-    return model.to(device)
+    return model.to(device)

diffrhythm/model/cfm.py

@@ -1,10 +1,22 @@
-"""
-ein notation:
-b - batch
-n - sequence
-nt - text sequence
-nw - raw wave length
-d - dimension
+# Copyright (c) 2025 ASLP-LAB
+#               2025 Ziqian Ning   (ningziqian@mail.nwpu.edu.cn)
+#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)
+#               2025 Guobin Ma     (guobin.ma@mail.nwpu.edu.cn)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""" This implementation is adapted from github repo:
+    https://github.com/SWivid/F5-TTS.
 """
 
 from __future__ import annotations
@@ -19,9 +31,7 @@ from torch.nn.utils.rnn import pad_sequence
 
 from torchdiffeq import odeint
 
-from diffrhythm.model.modules import MelSpec
 from diffrhythm.model.utils import (
-    default,
     exists,
     list_str_to_idx,
     list_str_to_tensor,
@@ -29,12 +39,25 @@ from diffrhythm.model.utils import (
     mask_from_frac_lengths,
 )
 
-def custom_mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"], device, max_seq_len):  # noqa: F722 F821
+def custom_mask_from_start_end_indices(
+    seq_len: int["b"],  # noqa: F821
+    latent_pred_segments,
+    device,
+    max_seq_len
+):
     max_seq_len = max_seq_len
     seq = torch.arange(max_seq_len, device=device).long()
-    start_mask = seq[None, :] >= start[:, None]
-    end_mask = seq[None, :] < end[:, None]
-    return start_mask & end_mask
+
+    res_mask = torch.zeros(max_seq_len, device=device, dtype=torch.bool)
+    
+    for start, end in latent_pred_segments:
+        start = start.unsqueeze(0)
+        end = end.unsqueeze(0)
+        start_mask = seq[None, :] >= start[:, None]
+        end_mask = seq[None, :] < end[:, None]
+        res_mask = res_mask | (start_mask & end_mask)
+    
+    return res_mask
 
 class CFM(nn.Module):
     def __init__(
@@ -42,7 +65,7 @@ class CFM(nn.Module):
         transformer: nn.Module,
         sigma=0.0,
         odeint_kwargs: dict = dict(
-            method="euler" # 'midpoint'
+            method="euler"
         ),
         odeint_options: dict = dict(
             min_step=0.05
@@ -54,7 +77,7 @@ class CFM(nn.Module):
         num_channels=None,
         frac_lengths_mask: tuple[float, float] = (0.7, 1.0),
         vocab_char_map: dict[str:int] | None = None,
-        use_style_prompt: bool = False
+        max_frames=2048
     ):
         super().__init__()
 
@@ -83,8 +106,8 @@ class CFM(nn.Module):
 
         # vocab map for tokenization
         self.vocab_char_map = vocab_char_map
-
-        self.use_style_prompt = use_style_prompt
+        
+        self.max_frames = max_frames
 
     @property
     def device(self):
@@ -112,10 +135,10 @@ class CFM(nn.Module):
         t_inter=0.1,
         edit_mask=None,
         start_time=None,
-        latent_pred_start_frame=0,
-        latent_pred_end_frame=2048,
+        latent_pred_segments=None,
         vocal_flag=False,
-        odeint_method="euler"
+        odeint_method="euler",
+        batch_infer_num=5
     ):
         self.eval()
         
@@ -125,7 +148,6 @@ class CFM(nn.Module):
             cond = cond.half()
 
         # raw wave
-        
         if cond.shape[1] > duration:
             cond = cond[:, :duration, :]
 
@@ -139,7 +161,6 @@ class CFM(nn.Module):
             lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)
 
         # text
-
         if isinstance(text, list):
             if exists(self.vocab_char_map):
                 text = list_str_to_idx(text, self.vocab_char_map).to(device)
@@ -147,26 +168,18 @@ class CFM(nn.Module):
                 text = list_str_to_tensor(text).to(device)
             assert text.shape[0] == batch
 
-        if exists(text):
-            text_lens = (text != -1).sum(dim=-1)
-
-
         # duration
         cond_mask = lens_to_mask(lens)
         if edit_mask is not None:
             cond_mask = cond_mask & edit_mask
 
-        latent_pred_start_frame = torch.tensor([latent_pred_start_frame]).to(cond.device)
-        latent_pred_end_frame = duration
-        latent_pred_end_frame = torch.tensor([latent_pred_end_frame]).to(cond.device)
-        fixed_span_mask = custom_mask_from_start_end_indices(cond_seq_len, latent_pred_start_frame, latent_pred_end_frame, device=cond.device, max_seq_len=duration)
-
+        latent_pred_segments = torch.tensor(latent_pred_segments).to(cond.device)
+        fixed_span_mask = custom_mask_from_start_end_indices(cond_seq_len, latent_pred_segments, device=cond.device, max_seq_len=duration)
         fixed_span_mask = fixed_span_mask.unsqueeze(-1)
         step_cond = torch.where(fixed_span_mask, torch.zeros_like(cond), cond)
 
         if isinstance(duration, int):
-            duration = torch.full((batch,), duration, device=device, dtype=torch.long)
-
+            duration = torch.full((batch_infer_num,), duration, device=device, dtype=torch.long)
 
         duration = duration.clamp(max=max_duration)
         max_duration = duration.amax()
@@ -175,7 +188,6 @@ class CFM(nn.Module):
         if duplicate_test:
             test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0)
 
-
         if batch > 1:
             mask = lens_to_mask(duration)
         else:  # save memory and speed up, as single inference need no mask currently
@@ -184,20 +196,27 @@ class CFM(nn.Module):
         # test for no ref audio
         if no_ref_audio:
             cond = torch.zeros_like(cond)
+            
+        if vocal_flag:
+            style_prompt = negative_style_prompt
+            negative_style_prompt = torch.zeros_like(style_prompt)
+
+        cond = cond.repeat(batch_infer_num, 1, 1)
+        step_cond = step_cond.repeat(batch_infer_num, 1, 1)
+        text = text.repeat(batch_infer_num, 1)
+        style_prompt = style_prompt.repeat(batch_infer_num, 1)
+        negative_style_prompt = negative_style_prompt.repeat(batch_infer_num, 1)
+        start_time = start_time.repeat(batch_infer_num)
+        fixed_span_mask = fixed_span_mask.repeat(batch_infer_num, 1, 1)
         
         start_time_embed, positive_text_embed, positive_text_residuals = self.transformer.forward_timestep_invariant(text, step_cond.shape[1], drop_text=False, start_time=start_time)
         _, negative_text_embed, negative_text_residuals = self.transformer.forward_timestep_invariant(text, step_cond.shape[1], drop_text=True, start_time=start_time)
 
-        if vocal_flag:
-            style_prompt = negative_style_prompt
-            negative_style_prompt = torch.zeros_like(style_prompt)
-            
         text_embed = torch.cat([positive_text_embed, negative_text_embed], 0)
         text_residuals = [torch.cat([a, b], 0) for a, b in zip(positive_text_residuals, negative_text_residuals)]
         step_cond = torch.cat([step_cond, step_cond], 0)
         style_prompt = torch.cat([style_prompt, negative_style_prompt], 0)
         start_time_embed = torch.cat([start_time_embed, start_time_embed], 0)
-            
 
         def fn(t, x):
             x = torch.cat([x, x], 0)
@@ -228,7 +247,7 @@ class CFM(nn.Module):
             t_start = t_inter
             y0 = (1 - t_start) * y0 + t_start * test_cond
             steps = int(steps * (1 - t_start))
-
+        
         t = torch.linspace(t_start, 1, steps, device=self.device, dtype=step_cond.dtype)
         if sway_sampling_coef is not None:
             t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
@@ -243,6 +262,7 @@ class CFM(nn.Module):
             out = out.permute(0, 2, 1)
             out = vocoder(out)
 
+        out = torch.chunk(out, batch_infer_num, dim=0)
         return out, trajectory
 
     def forward(
@@ -267,11 +287,10 @@ class CFM(nn.Module):
 
         # get a random span to mask out for training conditionally
         frac_lengths = torch.zeros((batch,), device=self.device).float().uniform_(*self.frac_lengths_mask)
-        rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)
+        rand_span_mask = mask_from_frac_lengths(lens, frac_lengths, self.max_frames)
 
         if exists(mask):
             rand_span_mask = mask
-            # rand_span_mask &= mask
 
         # mel is x1
         x1 = inp
@@ -301,7 +320,7 @@ class CFM(nn.Module):
         # adding mask will use more memory, thus also need to adjust batchsampler with scaled down threshold for long sequences
         pred = self.transformer(
             x=φ, cond=cond, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text, drop_prompt=drop_prompt,
-            style_prompt=style_prompt, style_prompt_lens=style_prompt_lens, grad_ckpt=grad_ckpt, start_time=start_time
+            style_prompt=style_prompt, start_time=start_time
         )
 
         # flow matching loss

diffrhythm/model/dit.py

@@ -1,10 +1,22 @@
-"""
-ein notation:
-b - batch
-n - sequence
-nt - text sequence
-nw - raw wave length
-d - dimension
+# Copyright (c) 2025 ASLP-LAB
+#               2025 Ziqian Ning   (ningziqian@mail.nwpu.edu.cn)
+#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)
+#               2025 Yuepeng Jiang (Jiangyp@mail.nwpu.edu.cn)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""" This implementation is adapted from github repo:
+    https://github.com/SWivid/F5-TTS.
 """
 
 from __future__ import annotations
@@ -12,22 +24,19 @@ from __future__ import annotations
 import torch
 from torch import nn
 import torch
-import torch.nn.functional as F
+
 from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRotaryEmbedding
 from transformers.models.llama import LlamaConfig
-from torch.utils.checkpoint import checkpoint
 
 from diffrhythm.model.modules import (
     TimestepEmbedding,
     ConvNeXtV2Block,
     ConvPositionEmbedding,
-    DiTBlock,
     AdaLayerNormZero_Final,
     precompute_freqs_cis,
     get_pos_embed_indices,
+    _prepare_decoder_attention_mask,
 )
-# from liger_kernel.transformers import apply_liger_kernel_to_llama
-# apply_liger_kernel_to_llama()
 
 # Text embedding
 class TextEmbedding(nn.Module):
@@ -77,7 +86,6 @@ class InputEmbedding(nn.Module):
     def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], style_emb, time_emb, drop_audio_cond=False):  # noqa: F722
         if drop_audio_cond:  # cfg for cond audio
             cond = torch.zeros_like(cond)
-
         style_emb = style_emb.unsqueeze(1).repeat(1, x.shape[1], 1)
         time_emb = time_emb.unsqueeze(1).repeat(1, x.shape[1], 1)
         x = self.proj(torch.cat((x, cond, text_embed, style_emb, time_emb), dim=-1))
@@ -85,9 +93,7 @@ class InputEmbedding(nn.Module):
         return x
 
 
-# Transformer backbone using DiT blocks
-
-
+# Transformer backbone using Llama blocks
 class DiT(nn.Module):
     def __init__(
         self,
@@ -103,26 +109,25 @@ class DiT(nn.Module):
         text_dim=None,
         conv_layers=0,
         long_skip_connection=False,
-        use_style_prompt=False,
-        max_pos=2048,
+        max_frames=2048
     ):
         super().__init__()
+        
+        self.max_frames = max_frames
 
         cond_dim = 512
         self.time_embed = TimestepEmbedding(cond_dim)
         self.start_time_embed = TimestepEmbedding(cond_dim)
         if text_dim is None:
             text_dim = mel_dim
-        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers, max_pos=max_pos)
+        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers, max_pos=self.max_frames)
         self.input_embed = InputEmbedding(mel_dim, text_dim, dim, cond_dim=cond_dim)
 
-
         self.dim = dim
         self.depth = depth
 
-        llama_config = LlamaConfig(hidden_size=dim, intermediate_size=dim * ff_mult, hidden_act='silu', max_position_embeddings=max_pos)
+        llama_config = LlamaConfig(hidden_size=dim, intermediate_size=dim * ff_mult, hidden_act='silu', max_position_embeddings=self.max_frames)
         llama_config._attn_implementation = 'sdpa'
-
         self.transformer_blocks = nn.ModuleList(
             [LlamaDecoderLayer(llama_config, layer_idx=i) for i in range(depth)]
         )
@@ -144,7 +149,6 @@ class DiT(nn.Module):
         self.norm_out = AdaLayerNormZero_Final(dim, cond_dim)  # final modulation
         self.proj_out = nn.Linear(dim, mel_dim)
 
-
     def forward_timestep_invariant(self, text, seq_len, drop_text, start_time):
         s_t = self.start_time_embed(start_time)
         text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
@@ -187,11 +191,22 @@ class DiT(nn.Module):
         pos_ids = torch.arange(x.shape[1], device=x.device)
         pos_ids = pos_ids.unsqueeze(0).repeat(x.shape[0], 1)
         rotary_embed = self.rotary_emb(x, pos_ids)
+        
+        attention_mask = torch.ones(
+            (batch, seq_len),
+            dtype=torch.bool,
+            device=x.device,
+        )
+        attention_mask = _prepare_decoder_attention_mask(
+            attention_mask,
+            (batch, seq_len),
+            x,
+        )
 
         for i, block in enumerate(self.transformer_blocks):
-            x, *_ = block(x, position_embeddings=rotary_embed)
+            x, *_ = block(x, attention_mask=attention_mask, position_embeddings=rotary_embed)
             if i < self.depth // 2:
-                x = x + text_residuals[i]
+                x = x + self.text_fusion_linears[i](text_embed)
 
         if self.long_skip_connection is not None:
             x = self.long_skip_connection(torch.cat((x, residual), dim=-1))

diffrhythm/model/modules.py

@@ -609,3 +609,44 @@ class TimestepEmbedding(nn.Module):
         time_hidden = time_hidden.to(timestep.dtype)
         time = self.time_mlp(time_hidden)  # b d
         return time
+
+
+# attention mask realated
+
+
+def _prepare_decoder_attention_mask(attention_mask, input_shape, inputs_embeds):
+    # create noncausal mask
+    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+    combined_attention_mask = None
+
+    def _expand_mask(
+        mask: torch.Tensor, dtype: torch.dtype, tgt_len: int = None
+    ):
+        """
+        Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+        """
+        bsz, src_len = mask.size()
+        tgt_len = tgt_len if tgt_len is not None else src_len
+
+        expanded_mask = (
+            mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+        )
+
+        inverted_mask = 1.0 - expanded_mask
+
+        return inverted_mask.masked_fill(
+            inverted_mask.to(torch.bool), torch.finfo(dtype).min
+        )
+
+    if attention_mask is not None:
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        expanded_attn_mask = _expand_mask(
+            attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+        ).to(inputs_embeds.device)
+        combined_attention_mask = (
+            expanded_attn_mask
+            if combined_attention_mask is None
+            else expanded_attn_mask + combined_attention_mask
+        )
+
+    return combined_attention_mask

diffrhythm/model/utils.py

@@ -44,15 +44,15 @@ def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]:  # noqa
     return seq[None, :] < t[:, None]
 
 
-def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"]):  # noqa: F722 F821
-    max_seq_len = 2048
+def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"], max_frames):  # noqa: F722 F821
+    max_seq_len = max_frames
     seq = torch.arange(max_seq_len, device=start.device).long()
     start_mask = seq[None, :] >= start[:, None]
     end_mask = seq[None, :] < end[:, None]
     return start_mask & end_mask
 
 
-def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"]):  # noqa: F722 F821
+def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"], max_frames):  # noqa: F722 F821
     lengths = (frac_lengths * seq_len).long()
     max_start = seq_len - lengths
 
@@ -60,7 +60,7 @@ def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"]):  # noqa
     start = (max_start * rand).long().clamp(min=0)
     end = start + lengths
 
-    return mask_from_start_end_indices(seq_len, start, end)
+    return mask_from_start_end_indices(seq_len, start, end, max_frames)
 
 
 def maybe_masked_mean(t: float["b n d"], mask: bool["b n"] = None) -> float["b d"]:  # noqa: F722

pretrained/eval.py

@@ -0,0 +1,66 @@
+from einops import rearrange
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class Generator(nn.Module):
+
+    def __init__(self,
+                 in_features,
+                 ffd_hidden_size,
+                 num_classes,
+                 attn_layer_num,
+                 
+                 ):
+        super(Generator, self).__init__()
+        
+        self.attn = nn.ModuleList(
+            [
+                nn.MultiheadAttention(
+                    embed_dim=in_features,
+                    num_heads=8,
+                    dropout=0.2,
+                    batch_first=True,
+                )
+                for _ in range(attn_layer_num)
+            ]
+        )
+        
+        self.ffd = nn.Sequential(
+            nn.Linear(in_features, ffd_hidden_size),
+            nn.ReLU(),
+            nn.Linear(ffd_hidden_size, in_features)
+        )
+        
+        self.dropout = nn.Dropout(0.2)
+        
+        self.fc =  nn.Linear(in_features * 2, num_classes)
+        
+        self.proj = nn.Tanh()
+        
+
+    def forward(self, ssl_feature, judge_id=None):
+        '''
+        ssl_feature: [B, T, D]   
+        output: [B, num_classes]
+        '''
+        
+        B, T, D = ssl_feature.shape
+        
+        ssl_feature = self.ffd(ssl_feature)
+        
+        tmp_ssl_feature = ssl_feature
+        
+        for attn in self.attn:
+            tmp_ssl_feature, _ = attn(tmp_ssl_feature, tmp_ssl_feature, tmp_ssl_feature)
+    
+        ssl_feature = self.dropout(torch.concat([torch.mean(tmp_ssl_feature, dim=1), torch.max(ssl_feature, dim=1)[0]], dim=1))  # B, 2D
+        
+        x = self.fc(ssl_feature)  # B, num_classes
+        
+        x = self.proj(x) * 2.0 + 3
+        
+        return x
+    
+    

pretrained/eval.safetensors

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

pretrained/eval.yaml

@@ -0,0 +1,6 @@
+generator:
+  _target_: pretrained.eval.Generator
+  in_features: 1024
+  ffd_hidden_size: 4096
+  num_classes: 5
+  attn_layer_num: 4