Added jam space

TangoFlux.py

@@ -1,58 +0,0 @@
-from diffusers import AutoencoderOobleck
-import torch
-from transformers  import T5EncoderModel,T5TokenizerFast
-from diffusers import   FluxTransformer2DModel
-from torch import nn
-from typing import List
-from diffusers import FlowMatchEulerDiscreteScheduler
-from diffusers.training_utils import compute_density_for_timestep_sampling
-import copy
-import torch.nn.functional as F
-import numpy as np
-from model import TangoFlux
-from huggingface_hub import snapshot_download
-from tqdm import tqdm
-from typing import Optional,Union,List
-from datasets import load_dataset, Audio
-from math import pi
-import json
-import inspect
-import yaml
-from safetensors.torch import load_file
-import os
-print(os.environ['HOME'])
-
-class TangoFluxInference:
-
-    def __init__(self,name='declare-lab/TangoFlux',device="cuda"):
-        
-
-        self.vae = AutoencoderOobleck.from_pretrained("stabilityai/stable-audio-open-1.0",subfolder='vae',token=os.environ['HF_TOKEN'])
-
-        paths = snapshot_download(repo_id=name,token=os.environ['HF_TOKEN']) 
-        weights = load_file("{}/model_1.safetensors".format(paths))
-
-        with open('{}/config.json'.format(paths),'r') as f:
-            config = json.load(f)
-        self.model = TangoFlux(config)
-        self.model.load_state_dict(weights,strict=False)
-        # _IncompatibleKeys(missing_keys=['text_encoder.encoder.embed_tokens.weight'], unexpected_keys=[]) this behaviour is expected 
-        self.vae.to(device)
-        self.model.to(device)
-    
-    def generate(self,prompt,steps=25,duration=10,guidance_scale=4.5):
-        
-        with torch.no_grad():
-            latents = self.model.inference_flow(prompt,
-            duration=duration,
-            num_inference_steps=steps,
-            guidance_scale=guidance_scale)
-
-
-
-            wave = self.vae.decode(latents.transpose(2,1)).sample.cpu()[0]
-        return wave
-
-
-
-

app.py

@@ -1,155 +1,68 @@
-import spaces
 import gradio as gr
-import json
-import torch
-import wavio
-from tqdm import tqdm
-from huggingface_hub import snapshot_download
-from pydub import AudioSegment
-from gradio import Markdown
-import uuid
-import torch
-from diffusers import DiffusionPipeline,AudioPipelineOutput
-from transformers import CLIPTextModel, T5EncoderModel, AutoModel, T5Tokenizer, T5TokenizerFast
-from typing import Union
-from diffusers.utils.torch_utils import randn_tensor
-from tqdm import tqdm
-from TangoFlux import TangoFluxInference
-import torchaudio
-
-
-
-tangoflux = TangoFluxInference(name="techneto/tangoflux-music-base-24k-120k-steps")
-
-
-
-@spaces.GPU(duration=15)
-def gradio_generate(prompt, steps, guidance,duration=10):
-
-    output = tangoflux.generate(prompt,steps=steps,guidance_scale=guidance,duration=duration)
-    #output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+import os
+from model import Jamify
+
+# Initialize the Jamify model once
+print("Initializing Jamify model...")
+jamify_model = Jamify()
+print("Jamify model ready.")
+
+def generate_song(reference_audio, lyrics_file, style_prompt, duration):
+    # We need to save the uploaded files to temporary paths to pass to the model
+    ref_audio_path = reference_audio.name if reference_audio else None
+    lyrics_path = lyrics_file.name
+
+    # The model expects paths, so we write the prompt to a temp file if needed
+    # (This part of the model could be improved to accept the string directly)
+
+    output_path = jamify_model.predict(
+        reference_audio_path=ref_audio_path,
+        lyrics_json_path=lyrics_path,
+        style_prompt=style_prompt,
+        duration_sec=duration
+    )
     
-    
-    #wavio.write(output_filename, output_wave, rate=44100, sampwidth=2)
-    filename = 'temp.wav'
-    #print(f"Saving audio to file: {unique_filename}")
-
-    # Save to file
-    torchaudio.save(filename, output, 44100)
-   
-
-    # Return the path to the generated audio file
-    return filename
-
-    #if (output_format == "mp3"):
-     #   AudioSegment.from_wav("temp.wav").export("temp.mp3", format = "mp3")
-      #  output_filename = "temp.mp3"
-
-    #return output_filename
-
-description_text = """
-Generate high quality and faithful audio in just a few seconds using <b>TangoFlux</b> by providing a text prompt. <b>TangoFlux</b> was trained from scratch and underwent alignment to follow human instructions using a new method called <b>Claped-Ranked Preference Optimization (CRPO)</b>.
-<div style="display: flex; gap: 10px; align-items: center;">
-  <a href="https://arxiv.org/abs/2412.21037">
-    <img src="https://img.shields.io/badge/Read_the_Paper-blue?link=https%3A%2F%2Fopenreview.net%2Fattachment%3Fid%3DtpJPlFTyxd%26name%3Dpdf" alt="arXiv">
-  </a>
-  <a href="https://huggingface.co/declare-lab/TangoFlux">
-    <img src="https://img.shields.io/badge/TangoFlux-Huggingface-violet?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fdeclare-lab%2FTangoFlux" alt="Static Badge">
-  </a>
-  <a href="https://tangoflux.github.io/">
-    <img src="https://img.shields.io/badge/Demos-declare--lab-brightred?style=flat" alt="Static Badge">
-  </a>
-  <a href="https://huggingface.co/spaces/declare-lab/TangoFlux">
-    <img src="https://img.shields.io/badge/TangoFlux-Huggingface_Space-8A2BE2?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fdeclare-lab%2FTangoFlux" alt="Static Badge">
-  </a>
-  <a href="https://huggingface.co/datasets/declare-lab/CRPO">
-    <img src="https://img.shields.io/badge/TangoFlux_Dataset-Huggingface-red?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fdatasets%2Fdeclare-lab%2FTangoFlux" alt="Static Badge">
-  </a>
-  <a href="https://github.com/declare-lab/TangoFlux">
-    <img src="https://img.shields.io/badge/Github-brown?logo=github&link=https%3A%2F%2Fgithub.com%2Fdeclare-lab%2FTangoFlux" alt="Static Badge">
-  </a>
-</div>
-"""
-# Gradio input and output components
-input_text = gr.Textbox(lines=2, label="Prompt")
-#output_format = gr.Radio(label = "Output format", info = "The file you can dowload", choices =  "wav"], value = "wav")
-output_audio = gr.Audio(label="Generated Audio", type="filepath")
-denoising_steps = gr.Slider(minimum=10, maximum=100, value=25, step=5, label="Steps", interactive=True)
-guidance_scale = gr.Slider(minimum=1, maximum=10, value=4.5, step=0.5, label="Guidance Scale", interactive=True)
-duration_scale = gr.Slider(minimum=1, maximum=30, value=10, step=1, label="Duration", interactive=True)
-
+    return output_path
 
 # Gradio interface
-gr_interface = gr.Interface(
-    fn=gradio_generate,
-    inputs=[input_text, denoising_steps, guidance_scale,duration_scale],
-    outputs=output_audio,
-    title="TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching and Clap-Ranked Preference Optimization",
-    description=description_text,
-    allow_flagging=False,
-    examples=[
-["A parade marches through a town square, with drumbeats pounding, children clapping, and a horse neighing amidst the commotion"],
-["A soccer ball hits a goalpost with a metallic clang, followed by cheers, clapping, and the distant hum of a commentator’s voice"],
-["The deep growl of an alligator ripples through the swamp as reeds sway with a soft rustle and a turtle splashes into the murky water"],        
-["A basketball bounces rhythmically on a court, shoes squeak against the floor, and a referee’s whistle cuts through the air"],
-["A train conductor blows a sharp whistle, metal wheels screech on the rails, and passengers murmur while settling into their seats"],
-["A fork scrapes a plate, water drips slowly into a sink, and the faint hum of a refrigerator lingers in the background"],
-["Alarms blare with rising urgency as fragments clatter against a metallic hull, interrupted by a faint hiss of escaping air"],
-["Tiny pops and hisses of chemical reactions intermingle with the rhythmic pumping of a centrifuge and the soft whirr of air filtration"],
-["A train conductor blows a sharp whistle, metal wheels screech on the rails, and passengers murmur while settling into their seats"],
-["Simulate a forest ambiance with birds chirping and wind rustling through the leaves"],
-["Quiet whispered conversation gradually fading into distant jet engine roar diminishing into silence"],
-["Clear sound of bicycle tires crunching on loose gravel and dirt, followed by deep male laughter echoing"],
-["Multiple ducks quacking loudly with splashing water and piercing wild animal shriek in background"],
-["Create a serene soundscape of a quiet beach at sunset"],
-["A pile of coins spills onto a wooden table with a metallic clatter, followed by the hushed murmur of a tavern crowd and the creak of a swinging door"],
-["Generate an energetic and bustling city street scene with distant traffic and close conversations"],
-["Powerful ocean waves crashing and receding on sandy beach with distant seagulls"],
-["Gentle female voice cooing and baby responding with happy gurgles and giggles"],
-["Clear male voice speaking, sharp popping sound, followed by genuine group laughter"],
-["Stream of water hitting empty ceramic cup, pitch rising as cup fills up"],
-["Massive crowd erupting in thunderous applause and excited cheering"],
-["Deep rolling thunder with bright lightning strikes crackling through sky"],
-["Aggressive dog barking and distressed cat meowing as racing car roars past at high speed"],
-["Peaceful stream bubbling and birds singing, interrupted by sudden explosive gunshot"],
-["Man speaking outdoors, goat bleating loudly, metal gate scraping closed, ducks quacking frantically, wind howling into microphone"],
-["Series of loud aggressive dog barks echoing"],
-["Multiple distinct cat meows at different pitches"],
-["Rhythmic wooden table tapping overlaid with steady water pouring sound"],
-["Sustained crowd applause with camera clicks and amplified male announcer voice"],
-["Two sharp gunshots followed by panicked birds taking flight with rapid wing flaps"],
-["Melodic human whistling harmonizing with natural birdsong"],
-["Deep rhythmic snoring with clear breathing patterns"],
-["Multiple racing engines revving and accelerating with sharp whistle piercing through"],
-["Massive stadium crowd cheering as thunder crashes and lightning strikes"],
-["Heavy helicopter blades chopping through air with engine and wind noise"],
-["Dog barking excitedly and man shouting as race car engine roars past"],
-["Quiet speech and then and airplane flying away"],
-["A bicycle peddling on dirt and gravel followed by a man speaking then laughing"],
-["Ducks quack and water splashes with some animal screeching in the background"],
-["Describe the sound of the ocean"],
-["A woman and a baby are having a conversation"],
-["A man speaks followed by a popping noise and laughter"],
-["A cup is filled from a faucet"],
-["An audience cheering and clapping"],
-["Rolling thunder with lightning strikes"],
-["A dog barking and a cat mewing and a racing car passes by"],
-["Gentle water stream, birds chirping and sudden gun shot"],
-["A dog barking"],
-["A cat meowing"],
-["Wooden table tapping sound while water pouring"],
-["Applause from a crowd with distant clicking and a man speaking over a loudspeaker"],
-["two gunshots followed by birds flying away while chirping"],
-["Whistling with birds chirping"],
-["A person snoring"],
-["Motor vehicles are driving with loud engines and a person whistles"],
-["People cheering in a stadium while thunder and lightning strikes"],
-["A helicopter is in flight"],
-["A dog barking and a man talking and a racing car passes by"],
-    ],
-    cache_examples="lazy", # Turn on to cache.
-)
-
-
-
-gr_interface.queue(15).launch()
+with gr.Blocks() as demo:
+    gr.Markdown("# Jamify: Music Generation from Lyrics and Style")
+    gr.Markdown("Provide your lyrics, a style reference (either an audio file or a text prompt), and a desired duration to generate a song.")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Inputs")
+            lyrics_file = gr.File(label="Lyrics File (.json)", type="filepath")
+            duration_slider = gr.Slider(minimum=5, maximum=180, value=30, step=1, label="Duration (seconds)")
+            
+            with gr.Tab("Style from Audio"):
+                reference_audio = gr.File(label="Reference Audio (.mp3, .wav)", type="filepath")
+            with gr.Tab("Style from Text"):
+                style_prompt = gr.Textbox(label="Style Prompt", lines=3, placeholder="e.g., A high-energy electronic dance track with a strong bassline and euphoric synths.")
+            
+            generate_button = gr.Button("Generate Song", variant="primary")
+            
+        with gr.Column():
+            gr.Markdown("### Output")
+            output_audio = gr.Audio(label="Generated Song")
+
+    generate_button.click(
+        fn=generate_song,
+        inputs=[reference_audio, lyrics_file, style_prompt, duration_slider],
+        outputs=output_audio,
+        api_name="generate_song"
+    )
+
+    gr.Markdown("### Example Usage")
+    gr.Examples(
+        examples=[
+            [None, "jamify/inputs/Jade Bird - Avalanche.json", "A sad, slow, acoustic country song", 30],
+            ["jamify/inputs/Rizzle Kicks, Rachel Chinouriri - Follow Excitement!.mp3", "jamify/inputs/Rizzle Kicks, Rachel Chinouriri - Follow Excitement!.json", "", 45],
+        ],
+        inputs=[reference_audio, lyrics_file, style_prompt, duration_slider],
+        outputs=output_audio,
+        fn=generate_song,
+        cache_examples=True
+    )
+
+demo.queue().launch()

gt0.json

@@ -0,0 +1 @@
+[{"word": "Every", "start_offset": 259, "end_offset": 267, "start": 20.72, "end": 21.36, "phoneme": "\u025bv\u025di|_"}, {"word": "night", "start_offset": 267, "end_offset": 275, "start": 21.36, "end": 22.0, "phoneme": "na\u026at|_"}, {"word": "in", "start_offset": 279, "end_offset": 283, "start": 22.32, "end": 22.64, "phoneme": "\u026an|_"}, {"word": "my", "start_offset": 283, "end_offset": 287, "start": 22.64, "end": 22.96, "phoneme": "ma\u026a|_"}, {"word": "dreams,", "start_offset": 287, "end_offset": 301, "start": 22.96, "end": 24.080000000000002, "phoneme": "dri\u02d0mz,"}, {"word": "I", "start_offset": 309, "end_offset": 313, "start": 24.72, "end": 25.04, "phoneme": "a\u026a|_"}, {"word": "see", "start_offset": 317, "end_offset": 321, "start": 25.36, "end": 25.68, "phoneme": "si\u02d0|_"}, {"word": "you,", "start_offset": 321, "end_offset": 325, "start": 25.68, "end": 26.0, "phoneme": "ju\u02d0,"}, {"word": "I", "start_offset": 340, "end_offset": 344, "start": 27.2, "end": 27.52, "phoneme": "a\u026a|_"}, {"word": "feel", "start_offset": 348, "end_offset": 352, "start": 27.84, "end": 28.16, "phoneme": "fi\u02d0l|_"}, {"word": "you.", "start_offset": 358, "end_offset": 362, "start": 28.64, "end": 28.96, "phoneme": "ju\u02d0."}, {"word": "That", "start_offset": 377, "end_offset": 381, "start": 30.16, "end": 30.48, "phoneme": "\u00f0\u00e6t|_"}, {"word": "is", "start_offset": 385, "end_offset": 389, "start": 30.8, "end": 31.12, "phoneme": "\u026az"}, {"word": "how", "start_offset": 393, "end_offset": 397, "start": 31.44, "end": 31.76, "phoneme": "ha\u028a|_"}, {"word": "I", "start_offset": 401, "end_offset": 405, "start": 32.08, "end": 32.4, "phoneme": "a\u026a|_"}, {"word": "know", "start_offset": 405, "end_offset": 409, "start": 32.4, "end": 32.72, "phoneme": "no\u028a|_"}, {"word": "you", "start_offset": 413, "end_offset": 417, "start": 33.04, "end": 33.36, "phoneme": "ju\u02d0|_"}, {"word": "go", "start_offset": 428, "end_offset": 431, "start": 34.24, "end": 34.480000000000004, "phoneme": "go\u028a|_"}, {"word": "far", "start_offset": 495, "end_offset": 503, "start": 39.6, "end": 40.24, "phoneme": "f\u0251\u02d0r"}, {"word": "across", "start_offset": 507, "end_offset": 517, "start": 40.56, "end": 41.36, "phoneme": "\u0259kr\u0254s|_"}, {"word": "the", "start_offset": 519, "end_offset": 523, "start": 41.52, "end": 41.84, "phoneme": "\u00f0\u0259|_"}, {"word": "distance", "start_offset": 527, "end_offset": 538, "start": 42.160000000000004, "end": 43.04, "phoneme": "d\u026ast\u0259ns|_"}, {"word": "and", "start_offset": 552, "end_offset": 556, "start": 44.160000000000004, "end": 44.480000000000004, "phoneme": "\u0259nd"}, {"word": "spaces", "start_offset": 556, "end_offset": 572, "start": 44.480000000000004, "end": 45.76, "phoneme": "spe\u026as\u0259z"}, {"word": "between", "start_offset": 583, "end_offset": 587, "start": 46.64, "end": 46.96, "phoneme": "b\u026atwi\u02d0n|_"}, {"word": "us.", "start_offset": 602, "end_offset": 606, "start": 48.160000000000004, "end": 48.480000000000004, "phoneme": "\u028cs."}, {"word": "You", "start_offset": 621, "end_offset": 625, "start": 49.68, "end": 50.0, "phoneme": "ju\u02d0|_"}, {"word": "have", "start_offset": 629, "end_offset": 633, "start": 50.32, "end": 50.64, "phoneme": "h\u00e6v"}, {"word": "come", "start_offset": 633, "end_offset": 637, "start": 50.64, "end": 50.96, "phoneme": "k\u028cm|_"}, {"word": "to", "start_offset": 641, "end_offset": 645, "start": 51.28, "end": 51.6, "phoneme": "tu\u02d0|_"}, {"word": "show", "start_offset": 649, "end_offset": 653, "start": 51.92, "end": 52.24, "phoneme": "\u0283o\u028a|_"}, {"word": "you", "start_offset": 655, "end_offset": 659, "start": 52.4, "end": 52.72, "phoneme": "ju\u02d0|_"}, {"word": "go", "start_offset": 673, "end_offset": 676, "start": 53.84, "end": 54.08, "phoneme": "go\u028a|_"}, {"word": "near,", "start_offset": 738, "end_offset": 745, "start": 59.04, "end": 59.6, "phoneme": "n\u026ar,"}, {"word": "far,", "start_offset": 768, "end_offset": 776, "start": 61.44, "end": 62.08, "phoneme": "f\u0251\u02d0r,"}, {"word": "wherever", "start_offset": 794, "end_offset": 806, "start": 63.52, "end": 64.48, "phoneme": "w\u025br\u025bv\u025d"}, {"word": "you", "start_offset": 822, "end_offset": 826, "start": 65.76, "end": 66.08, "phoneme": "ju\u02d0|_"}, {"word": "are.", "start_offset": 826, "end_offset": 830, "start": 66.08, "end": 66.4, "phoneme": "\u0251\u02d0r."}, {"word": "I", "start_offset": 849, "end_offset": 852, "start": 67.92, "end": 68.16, "phoneme": "a\u026a|_"}, {"word": "believe", "start_offset": 856, "end_offset": 868, "start": 68.48, "end": 69.44, "phoneme": "b\u026ali\u02d0v"}, {"word": "that", "start_offset": 875, "end_offset": 878, "start": 70.0, "end": 70.24, "phoneme": "\u00f0\u00e6t|_"}, {"word": "the", "start_offset": 886, "end_offset": 890, "start": 70.88, "end": 71.2, "phoneme": "\u00f0\u0259|_"}, {"word": "heart", "start_offset": 890, "end_offset": 898, "start": 71.2, "end": 71.84, "phoneme": "h\u0251\u02d0rt|_"}, {"word": "does", "start_offset": 898, "end_offset": 901, "start": 71.84, "end": 72.08, "phoneme": "d\u028cz"}, {"word": "go", "start_offset": 916, "end_offset": 920, "start": 73.28, "end": 73.60000000000001, "phoneme": "go\u028a|_"}, {"word": "on", "start_offset": 982, "end_offset": 985, "start": 78.56, "end": 78.8, "phoneme": "\u0251\u02d0n|_"}, {"word": "small.", "start_offset": 1009, "end_offset": 1017, "start": 80.72, "end": 81.36, "phoneme": "sm\u0254l."}, {"word": "You", "start_offset": 1037, "end_offset": 1041, "start": 82.96000000000001, "end": 83.28, "phoneme": "ju\u02d0|_"}, {"word": "open", "start_offset": 1045, "end_offset": 1049, "start": 83.60000000000001, "end": 83.92, "phoneme": "o\u028ap\u0259n|_"}, {"word": "the", "start_offset": 1065, "end_offset": 1069, "start": 85.2, "end": 85.52, "phoneme": "\u00f0\u0259|_"}, {"word": "door,", "start_offset": 1069, "end_offset": 1076, "start": 85.52, "end": 86.08, "phoneme": "d\u0254r,"}, {"word": "and", "start_offset": 1090, "end_offset": 1094, "start": 87.2, "end": 87.52, "phoneme": "\u0259nd"}, {"word": "you'll", "start_offset": 1094, "end_offset": 1100, "start": 87.52, "end": 88.0, "phoneme": "j\u028c\u028al|_"}, {"word": "hear", "start_offset": 1103, "end_offset": 1108, "start": 88.24, "end": 88.64, "phoneme": "hi\u02d0r"}, {"word": "in", "start_offset": 1119, "end_offset": 1122, "start": 89.52, "end": 89.76, "phoneme": "\u026an|_"}, {"word": "my", "start_offset": 1126, "end_offset": 1130, "start": 90.08, "end": 90.4, "phoneme": "ma\u026a|_"}, {"word": "heart,", "start_offset": 1130, "end_offset": 1138, "start": 90.4, "end": 91.04, "phoneme": "h\u0251\u02d0rt,"}, {"word": "and", "start_offset": 1141, "end_offset": 1145, "start": 91.28, "end": 91.60000000000001, "phoneme": "\u0259nd"}, {"word": "my", "start_offset": 1157, "end_offset": 1161, "start": 92.56, "end": 92.88, "phoneme": "ma\u026a|_"}, {"word": "heart", "start_offset": 1165, "end_offset": 1173, "start": 93.2, "end": 93.84, "phoneme": "h\u0251\u02d0rt|_"}, {"word": "will", "start_offset": 1173, "end_offset": 1177, "start": 93.84, "end": 94.16, "phoneme": "w\u026al|_"}, {"word": "go", "start_offset": 1185, "end_offset": 1189, "start": 94.8, "end": 95.12, "phoneme": "go\u028a|_"}, {"word": "and", "start_offset": 1211, "end_offset": 1215, "start": 96.88, "end": 97.2, "phoneme": "\u0259nd"}, {"word": "dawn.", "start_offset": 1223, "end_offset": 1233, "start": 97.84, "end": 98.64, "phoneme": "d\u0254n."}, {"word": "Love", "start_offset": 1345, "end_offset": 1353, "start": 107.60000000000001, "end": 108.24000000000001, "phoneme": "l\u028cv"}, {"word": "can", "start_offset": 1356, "end_offset": 1360, "start": 108.48, "end": 108.8, "phoneme": "k\u00e6n|_"}, {"word": "touch", "start_offset": 1360, "end_offset": 1366, "start": 108.8, "end": 109.28, "phoneme": "t\u028ct\u0283|_"}, {"word": "us", "start_offset": 1369, "end_offset": 1373, "start": 109.52, "end": 109.84, "phoneme": "\u028cs|_"}, {"word": "one", "start_offset": 1376, "end_offset": 1380, "start": 110.08, "end": 110.4, "phoneme": "w\u028cn|_"}, {"word": "time", "start_offset": 1384, "end_offset": 1388, "start": 110.72, "end": 111.04, "phoneme": "ta\u026am|_"}, {"word": "and", "start_offset": 1399, "end_offset": 1402, "start": 111.92, "end": 112.16, "phoneme": "\u0259nd"}, {"word": "last", "start_offset": 1406, "end_offset": 1410, "start": 112.48, "end": 112.8, "phoneme": "l\u00e6st|_"}, {"word": "for", "start_offset": 1416, "end_offset": 1420, "start": 113.28, "end": 113.60000000000001, "phoneme": "f\u0254r"}, {"word": "a", "start_offset": 1431, "end_offset": 1435, "start": 114.48, "end": 114.8, "phoneme": "\u0259|_"}, {"word": "lifetime", "start_offset": 1435, "end_offset": 1458, "start": 114.8, "end": 116.64, "phoneme": "la\u026afta\u026am|_"}, {"word": "and", "start_offset": 1471, "end_offset": 1475, "start": 117.68, "end": 118.0, "phoneme": "\u0259nd"}, {"word": "never", "start_offset": 1479, "end_offset": 1483, "start": 118.32000000000001, "end": 118.64, "phoneme": "n\u025bv\u025d"}, {"word": "let", "start_offset": 1487, "end_offset": 1491, "start": 118.96000000000001, "end": 119.28, "phoneme": "l\u025bt|_"}, {"word": "go", "start_offset": 1495, "end_offset": 1499, "start": 119.60000000000001, "end": 119.92, "phoneme": "go\u028a|_"}, {"word": "till", "start_offset": 1503, "end_offset": 1511, "start": 120.24000000000001, "end": 120.88, "phoneme": "t\u026al|_"}, {"word": "we're", "start_offset": 1521, "end_offset": 1528, "start": 121.68, "end": 122.24000000000001, "phoneme": "w\u025d\u02d0|_"}, {"word": "gone.", "start_offset": 1528, "end_offset": 1536, "start": 122.24000000000001, "end": 122.88, "phoneme": "g\u0254n."}, {"word": "Love", "start_offset": 1587, "end_offset": 1596, "start": 126.96000000000001, "end": 127.68, "phoneme": "l\u028cv"}, {"word": "was", "start_offset": 1599, "end_offset": 1603, "start": 127.92, "end": 128.24, "phoneme": "w\u0251\u02d0z"}, {"word": "when", "start_offset": 1607, "end_offset": 1611, "start": 128.56, "end": 128.88, "phoneme": "w\u025bn|_"}, {"word": "I", "start_offset": 1611, "end_offset": 1615, "start": 128.88, "end": 129.2, "phoneme": "a\u026a|_"}, {"word": "loved", "start_offset": 1615, "end_offset": 1626, "start": 129.2, "end": 130.08, "phoneme": "l\u028cvd"}, {"word": "you", "start_offset": 1626, "end_offset": 1630, "start": 130.08, "end": 130.4, "phoneme": "ju\u02d0|_"}, {"word": "one", "start_offset": 1641, "end_offset": 1644, "start": 131.28, "end": 131.52, "phoneme": "w\u028cn|_"}, {"word": "true", "start_offset": 1648, "end_offset": 1656, "start": 131.84, "end": 132.48, "phoneme": "tru\u02d0|_"}, {"word": "time.", "start_offset": 1656, "end_offset": 1660, "start": 132.48, "end": 132.8, "phoneme": "ta\u026am."}, {"word": "I", "start_offset": 1672, "end_offset": 1675, "start": 133.76, "end": 134.0, "phoneme": "a\u026a|_"}, {"word": "hold", "start_offset": 1679, "end_offset": 1687, "start": 134.32, "end": 134.96, "phoneme": "ho\u028ald"}, {"word": "to", "start_offset": 1691, "end_offset": 1693, "start": 135.28, "end": 135.44, "phoneme": "tu\u02d0|_"}, {"word": "in", "start_offset": 1712, "end_offset": 1716, "start": 136.96, "end": 137.28, "phoneme": "\u026an|_"}, {"word": "my", "start_offset": 1720, "end_offset": 1724, "start": 137.6, "end": 137.92000000000002, "phoneme": "ma\u026a|_"}, {"word": "life", "start_offset": 1724, "end_offset": 1728, "start": 137.92000000000002, "end": 138.24, "phoneme": "la\u026af|_"}, {"word": "will", "start_offset": 1731, "end_offset": 1733, "start": 138.48, "end": 138.64000000000001, "phoneme": "w\u026al|_"}, {"word": "always", "start_offset": 1743, "end_offset": 1747, "start": 139.44, "end": 139.76, "phoneme": "\u0254lwe\u026az"}, {"word": "go", "start_offset": 1763, "end_offset": 1767, "start": 141.04, "end": 141.36, "phoneme": "go\u028a|_"}, {"word": "near", "start_offset": 1830, "end_offset": 1836, "start": 146.4, "end": 146.88, "phoneme": "n\u026ar"}, {"word": "far", "start_offset": 1859, "end_offset": 1867, "start": 148.72, "end": 149.36, "phoneme": "f\u0251\u02d0r"}, {"word": "wherever", "start_offset": 1884, "end_offset": 1896, "start": 150.72, "end": 151.68, "phoneme": "w\u025br\u025bv\u025d"}, {"word": "you", "start_offset": 1914, "end_offset": 1918, "start": 153.12, "end": 153.44, "phoneme": "ju\u02d0|_"}, {"word": "are.", "start_offset": 1918, "end_offset": 1922, "start": 153.44, "end": 153.76, "phoneme": "\u0251\u02d0r."}, {"word": "I", "start_offset": 1940, "end_offset": 1943, "start": 155.20000000000002, "end": 155.44, "phoneme": "a\u026a|_"}, {"word": "believe", "start_offset": 1947, "end_offset": 1959, "start": 155.76, "end": 156.72, "phoneme": "b\u026ali\u02d0v"}, {"word": "that", "start_offset": 1966, "end_offset": 1970, "start": 157.28, "end": 157.6, "phoneme": "\u00f0\u00e6t|_"}, {"word": "the", "start_offset": 1974, "end_offset": 1977, "start": 157.92000000000002, "end": 158.16, "phoneme": "\u00f0\u0259|_"}, {"word": "heart", "start_offset": 1981, "end_offset": 1986, "start": 158.48, "end": 158.88, "phoneme": "h\u0251\u02d0rt|_"}, {"word": "does", "start_offset": 1990, "end_offset": 1993, "start": 159.20000000000002, "end": 159.44, "phoneme": "d\u028cz"}, {"word": "go", "start_offset": 2008, "end_offset": 2011, "start": 160.64000000000001, "end": 160.88, "phoneme": "go\u028a|_"}, {"word": "small.", "start_offset": 2099, "end_offset": 2111, "start": 167.92000000000002, "end": 168.88, "phoneme": "sm\u0254l."}, {"word": "You", "start_offset": 2127, "end_offset": 2131, "start": 170.16, "end": 170.48, "phoneme": "ju\u02d0|_"}, {"word": "open", "start_offset": 2136, "end_offset": 2140, "start": 170.88, "end": 171.20000000000002, "phoneme": "o\u028ap\u0259n|_"}, {"word": "the", "start_offset": 2156, "end_offset": 2160, "start": 172.48, "end": 172.8, "phoneme": "\u00f0\u0259|_"}, {"word": "door", "start_offset": 2160, "end_offset": 2167, "start": 172.8, "end": 173.36, "phoneme": "d\u0254r"}, {"word": "and", "start_offset": 2181, "end_offset": 2185, "start": 174.48, "end": 174.8, "phoneme": "\u0259nd"}, {"word": "you", "start_offset": 2185, "end_offset": 2187, "start": 174.8, "end": 174.96, "phoneme": "ju\u02d0|_"}, {"word": "hear", "start_offset": 2195, "end_offset": 2203, "start": 175.6, "end": 176.24, "phoneme": "hi\u02d0r"}, {"word": "in", "start_offset": 2209, "end_offset": 2213, "start": 176.72, "end": 177.04, "phoneme": "\u026an|_"}, {"word": "my", "start_offset": 2217, "end_offset": 2221, "start": 177.36, "end": 177.68, "phoneme": "ma\u026a|_"}, {"word": "heart,", "start_offset": 2221, "end_offset": 2230, "start": 177.68, "end": 178.4, "phoneme": "h\u0251\u02d0rt,"}, {"word": "and", "start_offset": 2232, "end_offset": 2236, "start": 178.56, "end": 178.88, "phoneme": "\u0259nd"}, {"word": "my", "start_offset": 2248, "end_offset": 2251, "start": 179.84, "end": 180.08, "phoneme": "ma\u026a|_"}, {"word": "heart", "start_offset": 2255, "end_offset": 2263, "start": 180.4, "end": 181.04, "phoneme": "h\u0251\u02d0rt|_"}, {"word": "will", "start_offset": 2263, "end_offset": 2266, "start": 181.04, "end": 181.28, "phoneme": "w\u026al|_"}, {"word": "go", "start_offset": 2278, "end_offset": 2282, "start": 182.24, "end": 182.56, "phoneme": "go\u028a|_"}, {"word": "on.", "start_offset": 2286, "end_offset": 2289, "start": 182.88, "end": 183.12, "phoneme": "\u0251\u02d0n."}, {"word": "You", "start_offset": 2557, "end_offset": 2559, "start": 204.56, "end": 204.72, "phoneme": "ju\u02d0|_"}, {"word": "hear", "start_offset": 2587, "end_offset": 2594, "start": 206.96, "end": 207.52, "phoneme": "hi\u02d0r"}, {"word": "there's", "start_offset": 2610, "end_offset": 2620, "start": 208.8, "end": 209.6, "phoneme": "\u00f0\u025brz"}, {"word": "nothing", "start_offset": 2620, "end_offset": 2632, "start": 209.6, "end": 210.56, "phoneme": "n\u028c\u03b8\u026a\u014b|_"}, {"word": "I", "start_offset": 2640, "end_offset": 2644, "start": 211.20000000000002, "end": 211.52, "phoneme": "a\u026a|_"}, {"word": "fear,", "start_offset": 2644, "end_offset": 2651, "start": 211.52, "end": 212.08, "phoneme": "f\u026ar,"}, {"word": "and", "start_offset": 2666, "end_offset": 2669, "start": 213.28, "end": 213.52, "phoneme": "\u0259nd"}, {"word": "I", "start_offset": 2673, "end_offset": 2677, "start": 213.84, "end": 214.16, "phoneme": "a\u026a|_"}, {"word": "know", "start_offset": 2677, "end_offset": 2681, "start": 214.16, "end": 214.48000000000002, "phoneme": "no\u028a|_"}, {"word": "that", "start_offset": 2693, "end_offset": 2697, "start": 215.44, "end": 215.76, "phoneme": "\u00f0\u00e6t|_"}, {"word": "my", "start_offset": 2701, "end_offset": 2705, "start": 216.08, "end": 216.4, "phoneme": "ma\u026a|_"}, {"word": "heart", "start_offset": 2705, "end_offset": 2713, "start": 216.4, "end": 217.04, "phoneme": "h\u0251\u02d0rt|_"}, {"word": "will", "start_offset": 2717, "end_offset": 2721, "start": 217.36, "end": 217.68, "phoneme": "w\u026al|_"}, {"word": "go", "start_offset": 2733, "end_offset": 2736, "start": 218.64000000000001, "end": 218.88, "phoneme": "go\u028a|_"}, {"word": "forever", "start_offset": 2852, "end_offset": 2863, "start": 228.16, "end": 229.04, "phoneme": "f\u025d\u025bv\u025d"}, {"word": "this", "start_offset": 2881, "end_offset": 2883, "start": 230.48000000000002, "end": 230.64000000000001, "phoneme": "\u00f0\u026as|_"}, {"word": "way.", "start_offset": 2888, "end_offset": 2892, "start": 231.04, "end": 231.36, "phoneme": "we\u026a."}, {"word": "You", "start_offset": 2908, "end_offset": 2911, "start": 232.64000000000001, "end": 232.88, "phoneme": "ju\u02d0|_"}, {"word": "are", "start_offset": 2914, "end_offset": 2918, "start": 233.12, "end": 233.44, "phoneme": "\u0251\u02d0r"}, {"word": "safe", "start_offset": 2928, "end_offset": 2935, "start": 234.24, "end": 234.8, "phoneme": "se\u026af|_"}, {"word": "in", "start_offset": 2938, "end_offset": 2942, "start": 235.04, "end": 235.36, "phoneme": "\u026an|_"}, {"word": "my", "start_offset": 2942, "end_offset": 2946, "start": 235.36, "end": 235.68, "phoneme": "ma\u026a|_"}, {"word": "heart,", "start_offset": 2950, "end_offset": 2957, "start": 236.0, "end": 236.56, "phoneme": "h\u0251\u02d0rt,"}, {"word": "and", "start_offset": 2959, "end_offset": 2963, "start": 236.72, "end": 237.04, "phoneme": "\u0259nd"}, {"word": "my", "start_offset": 2975, "end_offset": 2978, "start": 238.0, "end": 238.24, "phoneme": "ma\u026a|_"}, {"word": "heart", "start_offset": 2982, "end_offset": 2990, "start": 238.56, "end": 239.20000000000002, "phoneme": "h\u0251\u02d0rt|_"}, {"word": "will", "start_offset": 2990, "end_offset": 2994, "start": 239.20000000000002, "end": 239.52, "phoneme": "w\u026al|_"}, {"word": "go", "start_offset": 3002, "end_offset": 3005, "start": 240.16, "end": 240.4, "phoneme": "go\u028a|_"}, {"word": "on", "start_offset": 3009, "end_offset": 3012, "start": 240.72, "end": 240.96, "phoneme": "\u0251\u02d0n|_"}, {"word": "there.", "start_offset": 3028, "end_offset": 3032, "start": 242.24, "end": 242.56, "phoneme": "\u00f0\u025br."}]

model.py

@@ -1,511 +1,180 @@
-from transformers  import T5EncoderModel,T5TokenizerFast
-import torch
-from diffusers import   FluxTransformer2DModel
-from torch import nn
 
-from typing import List
-from diffusers import FlowMatchEulerDiscreteScheduler
-from diffusers.training_utils import compute_density_for_timestep_sampling
-import copy
-import torch.nn.functional as F
+import torch
+import torchaudio
+from omegaconf import OmegaConf
+from huggingface_hub import snapshot_download
 import numpy as np
-from tqdm import tqdm
-
-from typing import Optional,Union,List
-from datasets import load_dataset, Audio
-from math import pi
-import inspect
-import yaml
-
-
-
-class StableAudioPositionalEmbedding(nn.Module):
-    """Used for continuous time
-
-    Adapted from stable audio open.
-    
-    """
-
-    def __init__(self, dim: int):
-        super().__init__()
-        assert (dim % 2) == 0
-        half_dim = dim // 2
-        self.weights = nn.Parameter(torch.randn(half_dim))
-
-    def forward(self, times: torch.Tensor) -> torch.Tensor:
-        times = times[..., None]
-        freqs = times * self.weights[None] * 2 * pi
-        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
-        fouriered = torch.cat((times, fouriered), dim=-1)
-        return fouriered
-        
-class DurationEmbedder(nn.Module):
-    """
-    A simple linear projection model to map numbers to a latent space.
-
-    Code is adapted from
-    https://github.com/Stability-AI/stable-audio-tools
-
-    Args:
-        number_embedding_dim (`int`):
-            Dimensionality of the number embeddings.
-        min_value (`int`):
-            The minimum value of the seconds number conditioning modules.
-        max_value (`int`):
-            The maximum value of the seconds number conditioning modules
-        internal_dim (`int`):
-            Dimensionality of the intermediate number hidden states.
-    """
-
-    def __init__(
-        self,
-        number_embedding_dim,
-        min_value,
-        max_value,
-        internal_dim: Optional[int] = 256,
-    ):
-        super().__init__()
-        self.time_positional_embedding = nn.Sequential(
-            StableAudioPositionalEmbedding(internal_dim),
-            nn.Linear(in_features=internal_dim + 1, out_features=number_embedding_dim),
-        )
-
-        self.number_embedding_dim = number_embedding_dim
-        self.min_value = min_value
-        self.max_value = max_value
-        self.dtype = torch.float32 
-
-    def forward(
-        self,
-        floats: torch.Tensor,
-    ):
-        floats = floats.clamp(self.min_value, self.max_value)
-
-        normalized_floats = (floats - self.min_value) / (self.max_value - self.min_value)
-
-        # Cast floats to same type as embedder
-        embedder_dtype = next(self.time_positional_embedding.parameters()).dtype
-        normalized_floats = normalized_floats.to(embedder_dtype)
-
-        embedding = self.time_positional_embedding(normalized_floats)
-        float_embeds = embedding.view(-1, 1, self.number_embedding_dim)
-
-        return float_embeds
-
-
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-    
-
-
-       
-
-
-
-class TangoFlux(nn.Module):
-
-
-    def __init__(self,config,initialize_reference_model=False):
-
-        super().__init__()
-        
-
-        
-        self.num_layers = config.get('num_layers', 6)
-        self.num_single_layers = config.get('num_single_layers', 18)
-        self.in_channels = config.get('in_channels', 64)
-        self.attention_head_dim = config.get('attention_head_dim', 128)
-        self.joint_attention_dim = config.get('joint_attention_dim', 1024)
-        self.num_attention_heads = config.get('num_attention_heads', 8)
-        self.audio_seq_len = config.get('audio_seq_len', 645)
-        self.max_duration = config.get('max_duration', 30)
-        self.uncondition = config.get('uncondition', False)
-        self.text_encoder_name = config.get('text_encoder_name', "google/flan-t5-large")
-        
-        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000)
-        self.noise_scheduler_copy = copy.deepcopy(self.noise_scheduler)
-        self.max_text_seq_len = 64
-        self.text_encoder = T5EncoderModel.from_pretrained(self.text_encoder_name)
-        self.tokenizer = T5TokenizerFast.from_pretrained(self.text_encoder_name)
-        self.text_embedding_dim = self.text_encoder.config.d_model
-        
-        
-        self.fc = nn.Sequential(nn.Linear(self.text_embedding_dim,self.joint_attention_dim),nn.ReLU())
-        self.duration_emebdder = DurationEmbedder(self.text_embedding_dim,min_value=0,max_value=self.max_duration)
-        
-        self.transformer = FluxTransformer2DModel(
-                                     in_channels=self.in_channels,
-                                     num_layers=self.num_layers,
-                                     num_single_layers=self.num_single_layers,
-                                     attention_head_dim=self.attention_head_dim,
-                                     num_attention_heads=self.num_attention_heads,
-                                     joint_attention_dim=self.joint_attention_dim,
-                                     pooled_projection_dim=self.text_embedding_dim,
-                                     guidance_embeds=False)
-
-        self.beta_dpo = 2000 ## this is used for dpo training
-            
-            
-        
-        
-       
-        
-    def get_sigmas(self,timesteps, n_dim=3, dtype=torch.float32):
-        device = self.text_encoder.device
-        sigmas = self.noise_scheduler_copy.sigmas.to(device=device, dtype=dtype)
-        
-
-        schedule_timesteps = self.noise_scheduler_copy.timesteps.to(device)
-        timesteps = timesteps.to(device)
-        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
-       
-        sigma = sigmas[step_indices].flatten()
-        while len(sigma.shape) < n_dim:
-            sigma = sigma.unsqueeze(-1)
-        return sigma
-    
-   
-    
-    def encode_text_classifier_free(self, prompt: List[str], num_samples_per_prompt=1):
-        device = self.text_encoder.device
-        batch = self.tokenizer(
-            prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
-        )
-        input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)
-
-        with torch.no_grad():
-            prompt_embeds = self.text_encoder(
-                input_ids=input_ids, attention_mask=attention_mask
-            )[0]
-                
-        prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
-        attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)
-
-        # get unconditional embeddings for classifier free guidance
-        uncond_tokens = [""] 
-        
-        max_length = prompt_embeds.shape[1]
-        uncond_batch = self.tokenizer(
-            uncond_tokens, max_length=max_length, padding='max_length', truncation=True, return_tensors="pt",
-        )
-        uncond_input_ids = uncond_batch.input_ids.to(device)
-        uncond_attention_mask = uncond_batch.attention_mask.to(device)
-       
-        with torch.no_grad():
-            negative_prompt_embeds = self.text_encoder(
-                input_ids=uncond_input_ids, attention_mask=uncond_attention_mask
-            )[0]
-                
-        negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
-        uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)
-
-        # For classifier free guidance, we need to do two forward passes.
-        # We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes
-       
-        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
-        prompt_mask = torch.cat([uncond_attention_mask, attention_mask])
-        boolean_prompt_mask = (prompt_mask == 1).to(device)
-
-        return prompt_embeds, boolean_prompt_mask
-
-    @torch.no_grad()
-    def encode_text(self, prompt):
-        device = self.text_encoder.device
-        batch = self.tokenizer(
-            prompt, max_length=self.max_text_seq_len, padding=True, truncation=True, return_tensors="pt")
-        input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)
-
-        
-       
-        encoder_hidden_states = self.text_encoder(
-            input_ids=input_ids, attention_mask=attention_mask)[0]
-    
-        boolean_encoder_mask = (attention_mask == 1).to(device)
-        
-        return encoder_hidden_states, boolean_encoder_mask
-    
-        
-    def encode_duration(self,duration):
-        return self.duration_emebdder(duration)
-
-
-    
-    @torch.no_grad()
-    def inference_flow(self, prompt,
-                    num_inference_steps=50,
-                    timesteps=None,
-                    guidance_scale=3,
-                    duration=10,
-                    disable_progress=False,
-                    num_samples_per_prompt=1):
-
-        '''Only tested for single inference. Haven't test for batch inference'''
-
-        bsz = num_samples_per_prompt
-        device = self.transformer.device
-        scheduler = self.noise_scheduler
-
-        if not isinstance(prompt,list):
-            prompt = [prompt]
-        if not isinstance(duration,torch.Tensor):
-            duration = torch.tensor([duration],device=device)
-        classifier_free_guidance = guidance_scale > 1.0
-        duration_hidden_states = self.encode_duration(duration)
-        if classifier_free_guidance:
-            bsz = 2 * num_samples_per_prompt
-
-            encoder_hidden_states, boolean_encoder_mask = self.encode_text_classifier_free(prompt, num_samples_per_prompt=num_samples_per_prompt)
-            duration_hidden_states = duration_hidden_states.repeat(bsz,1,1)
-        
-
+import json
+import os
+from safetensors.torch import load_file
+
+# Imports from the jamify library
+from jam.model.cfm import CFM
+from jam.model.dit import DiT
+from jam.model.vae import StableAudioOpenVAE
+from jam.dataset import DiffusionWebDataset, enhance_webdataset_config
+from muq import MuQMuLan
+
+# Helper functions adapted from jamify/src/jam/infer.py
+def get_negative_style_prompt(device, file_path):
+    vocal_style = np.load(file_path)
+    vocal_style = torch.from_numpy(vocal_style).to(device)
+    return vocal_style.half()
+
+def normalize_audio(audio):
+    audio = audio - audio.mean(-1, keepdim=True)
+    audio = audio / (audio.abs().max(-1, keepdim=True).values + 1e-8)
+    return audio
+
+class Jamify:
+    def __init__(self, device="cuda" if torch.cuda.is_available() else "cpu"):
+        self.device = torch.device(device)
+        
+        # --- FIX: Point to the local jamify repository for config and public files ---
+        #jamify_repo_path = "/Users/cy/Desktop/JAM/jamify"
+        
+        print("Downloading main model checkpoint...")
+        model_repo_path = snapshot_download(repo_id="declare-lab/jam-0.5")
+        self.checkpoint_path = os.path.join(model_repo_path, "jam-0_5.safetensors")
+        
+        # Use local config and data files
+        config_path = os.path.join(model_repo_path, "jam_infer.yaml")
+        self.negative_style_prompt_path = os.path.join(model_repo_path, "vocal.npy")
+        tokenizer_path = os.path.join(model_repo_path, "en_us_cmudict_ipa_forward.pt")
+        silence_latent_path = os.path.join(model_repo_path, "silience_latent.pt")
+        print("Loading configuration...")
+        self.config = OmegaConf.load(config_path)
+        self.config.data.train_dataset.silence_latent_path = silence_latent_path
+        
+        # --- FIX: Override the relative paths in the config with absolute paths ---
+        self.config.data.train_dataset.tokenizer_path = tokenizer_path
+        self.config.evaluation.dataset.tokenizer_path = tokenizer_path
+        self.config.data.train_dataset.phonemizer_checkpoint = tokenizer_path
+        
+        print("Loading VAE model...")
+        self.vae = StableAudioOpenVAE().to(self.device).eval()
+        
+        print("Loading CFM model...")
+        self.cfm_model = self._load_cfm_model(self.config.model, self.checkpoint_path)
+        
+        print("Loading MuQ style model...")
+        self.muq_model = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large").to(self.device).eval()
+
+        print("Setting up dataset processor...")
+        dataset_cfg = OmegaConf.merge(self.config.data.train_dataset, self.config.evaluation.dataset)
+        enhance_webdataset_config(dataset_cfg)
+        dataset_cfg.multiple_styles = False
+        self.dataset_processor = DiffusionWebDataset(**dataset_cfg)
+
+        print("Jamify model loaded successfully.")
+
+    def _load_cfm_model(self, model_config, checkpoint_path):
+        dit_config = model_config["dit"].copy()
+        if "text_num_embeds" not in dit_config:
+            dit_config["text_num_embeds"] = 256
+        
+        model = CFM(
+            transformer=DiT(**dit_config),
+            **model_config["cfm"]
+        ).to(self.device)
+        
+        state_dict = load_file(checkpoint_path)
+        model.load_state_dict(state_dict, strict=False)
+        return model.eval()
+
+    def _generate_style_embedding_from_audio(self, audio_path):
+        waveform, sample_rate = torchaudio.load(audio_path)
+        if sample_rate != 24000:
+            resampler = torchaudio.transforms.Resample(sample_rate, 24000)
+            waveform = resampler(waveform)
+        if waveform.shape[0] > 1:
+            waveform = waveform.mean(dim=0, keepdim=True)
+        
+        waveform = waveform.squeeze(0).to(self.device)
+        
+        with torch.inference_mode():
+            style_embedding = self.muq_model(wavs=waveform.unsqueeze(0)[..., :24000 * 30])
+        return style_embedding[0]
+
+    def _generate_style_embedding_from_prompt(self, prompt):
+        with torch.inference_mode():
+            style_embedding = self.muq_model(texts=[prompt]).squeeze(0)
+        return style_embedding
+
+    def predict(self, reference_audio_path, lyrics_json_path, style_prompt, duration_sec=30, steps=50):
+        print("Starting prediction...")
+        
+        if reference_audio_path:
+            print(f"Generating style from audio: {reference_audio_path}")
+            style_embedding = self._generate_style_embedding_from_audio(reference_audio_path)
+        elif style_prompt:
+            print(f"Generating style from prompt: '{style_prompt}'")
+            style_embedding = self._generate_style_embedding_from_prompt(style_prompt)
         else:
+            print("No style provided, using zero embedding.")
+            style_embedding = torch.zeros(512, device=self.device)
 
-            encoder_hidden_states, boolean_encoder_mask = self.encode_text(prompt,num_samples_per_prompt=num_samples_per_prompt)
-            
-        mask_expanded = boolean_encoder_mask.unsqueeze(-1).expand_as(encoder_hidden_states)
-        masked_data = torch.where(mask_expanded, encoder_hidden_states, torch.tensor(float('nan')))
-
-        pooled = torch.nanmean(masked_data, dim=1)
-        pooled_projection = self.fc(pooled)
-
-        encoder_hidden_states = torch.cat([encoder_hidden_states,duration_hidden_states],dim=1) ## (bs,seq_len,dim)
-        
-        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
-        timesteps, num_inference_steps = retrieve_timesteps(
-            scheduler,
-            num_inference_steps,
-            device,
-            timesteps,
-            sigmas
-        )
-
-        latents = torch.randn(num_samples_per_prompt,self.audio_seq_len,64)
-        weight_dtype = latents.dtype
-
-        progress_bar = tqdm(range(num_inference_steps), disable=disable_progress)
-
-        txt_ids = torch.zeros(bsz,encoder_hidden_states.shape[1],3).to(device)
-        audio_ids = torch.arange(self.audio_seq_len).unsqueeze(0).unsqueeze(-1).repeat(bsz,1,3).to(device)
+        print(f"Loading lyrics from: {lyrics_json_path}")
+        with open(lyrics_json_path, 'r') as f:
+            lrc_data = json.load(f)
+        if 'word' not in lrc_data:
+            lrc_data = {'word': lrc_data}
 
+        frame_rate = 21.5
+        num_frames = int(duration_sec * frame_rate)
+        fake_latent = torch.randn(128, num_frames)
         
-        timesteps = timesteps.to(device)
-        latents = latents.to(device)
-        encoder_hidden_states = encoder_hidden_states.to(device)
+        sample_tuple = ("user_song", fake_latent, style_embedding, lrc_data)
         
+        print("Processing sample...")
+        processed_sample = self.dataset_processor.process_sample_safely(sample_tuple)
+        if processed_sample is None:
+            raise ValueError("Failed to process the provided lyrics and style.")
 
-        for i, t in enumerate(timesteps):
-            
-            latents_input = torch.cat([latents] * 2) if classifier_free_guidance else latents
-
+        batch = self.dataset_processor.custom_collate_fn([processed_sample])
         
+        for key, value in batch.items():
+            if isinstance(value, torch.Tensor):
+                batch[key] = value.to(self.device)
 
-            noise_pred = self.transformer(
-                    hidden_states=latents_input,
-                    # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
-                    timestep=torch.tensor([t/1000],device=device),
-                    guidance = None,
-                    pooled_projections=pooled_projection,
-                    encoder_hidden_states=encoder_hidden_states,
-                    txt_ids=txt_ids,
-                    img_ids=audio_ids,
-                    return_dict=False,
-                )[0]
+        print("Generating audio latent...")
+        with torch.inference_mode():
+            batch_size = 1
+            text = batch["lrc"]
+            style_prompt_tensor = batch["prompt"]
+            start_time = batch["start_time"]
+            duration_abs = batch["duration_abs"]
+            duration_rel = batch["duration_rel"]
             
-            if classifier_free_guidance:
-                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+            cond = torch.zeros(batch_size, self.cfm_model.max_frames, 64).to(self.device)
+            pred_frames = [(0, self.cfm_model.max_frames)]
             
-            
-            latents = scheduler.step(noise_pred, t, latents).prev_sample
-
-
-        return latents
-
-    def forward(self,
-                latents,
-                prompt,
-                duration=torch.tensor([10]),
-                sft=True
-                ):
-
+            negative_style_prompt = get_negative_style_prompt(self.device, self.negative_style_prompt_path)
+            negative_style_prompt = negative_style_prompt.repeat(batch_size, 1)
 
-        device = latents.device
-        audio_seq_length = self.audio_seq_len
-        bsz = latents.shape[0]
-        
-
-
-        encoder_hidden_states, boolean_encoder_mask = self.encode_text(prompt)
-        duration_hidden_states = self.encode_duration(duration)
-            
+            sample_kwargs = self.config.evaluation.sample_kwargs
+            sample_kwargs.steps = steps
+            latents, _ = self.cfm_model.sample(
+                cond=cond, text=text, style_prompt=style_prompt_tensor,
+                duration_abs=duration_abs, duration_rel=duration_rel,
+                negative_style_prompt=negative_style_prompt, start_time=start_time,
+                latent_pred_segments=pred_frames, **sample_kwargs)
             
-        mask_expanded = boolean_encoder_mask.unsqueeze(-1).expand_as(encoder_hidden_states)
-        masked_data = torch.where(mask_expanded, encoder_hidden_states, torch.tensor(float('nan')))
-        pooled = torch.nanmean(masked_data, dim=1)
-        pooled_projection = self.fc(pooled)
+            latent = latents[0][0]
 
-        ## Add duration hidden states to encoder hidden states
-        encoder_hidden_states = torch.cat([encoder_hidden_states,duration_hidden_states],dim=1) ## (bs,seq_len,dim)
-
-        txt_ids = torch.zeros(bsz,encoder_hidden_states.shape[1],3).to(device)
-        audio_ids = torch.arange(audio_seq_length).unsqueeze(0).unsqueeze(-1).repeat(bsz,1,3).to(device)
+        print("Decoding latent to audio...")
+        latent_for_vae = latent.transpose(0, 1).unsqueeze(0)
+        pred_audio = self.vae.decode(latent_for_vae).sample.squeeze(0).detach().cpu()
         
-        if sft:
-            
-            if self.uncondition:
-                mask_indices = [k for k in range(len(prompt)) if random.random() < 0.1]
-                if len(mask_indices) > 0:
-                    encoder_hidden_states[mask_indices] = 0
-            
-            
-            noise = torch.randn_like(latents)
-            
-            
-            u = compute_density_for_timestep_sampling(
-                    weighting_scheme='logit_normal',
-                    batch_size=bsz,
-                    logit_mean=0,
-                    logit_std=1,
-                    mode_scale=None,
-                )
-
-                
-            indices = (u * self.noise_scheduler_copy.config.num_train_timesteps).long()
-            timesteps = self.noise_scheduler_copy.timesteps[indices].to(device=latents.device)
-            sigmas = self.get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
-            
-            noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
-            
-            
-
-            model_pred =  self.transformer(
-                                    hidden_states=noisy_model_input,
-                                    encoder_hidden_states=encoder_hidden_states,
-                                    pooled_projections=pooled_projection,
-                                    img_ids=audio_ids,
-                                    txt_ids=txt_ids,
-                                    guidance=None,
-                # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
-                                    timestep=timesteps/1000,
-                                    return_dict=False)[0]
-            
-            
-
-            target = noise - latents
-            loss = torch.mean(
-                        ( (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
-                        1,
-                    )
-            loss = loss.mean()
-            raw_model_loss, raw_ref_loss,implicit_acc,epsilon_diff = 0,0,0,0 ## default this to 0 if doing sft
-
-        else:
-            encoder_hidden_states = encoder_hidden_states.repeat(2, 1, 1)
-            pooled_projection = pooled_projection.repeat(2,1)
-            noise = torch.randn_like(latents).chunk(2)[0].repeat(2, 1, 1) ## Have to sample same noise for preferred and rejected
-            u = compute_density_for_timestep_sampling(
-                    weighting_scheme='logit_normal',
-                    batch_size=bsz//2,
-                    logit_mean=0,
-                    logit_std=1,
-                    mode_scale=None,
-                )
-
-                
-            indices = (u * self.noise_scheduler_copy.config.num_train_timesteps).long()
-            timesteps = self.noise_scheduler_copy.timesteps[indices].to(device=latents.device)
-            timesteps = timesteps.repeat(2)
-            sigmas = self.get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
-            
-            noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
-
-            model_pred =  self.transformer(
-                                    hidden_states=noisy_model_input,
-                                    encoder_hidden_states=encoder_hidden_states,
-                                    pooled_projections=pooled_projection,
-                                    img_ids=audio_ids,
-                                    txt_ids=txt_ids,
-                                    guidance=None,
-                # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
-                                    timestep=timesteps/1000,
-                                    return_dict=False)[0]
-            target = noise - latents
-            
-            model_losses = F.mse_loss(model_pred.float(), target.float(), reduction="none")
-            model_losses = model_losses.mean(dim=list(range(1, len(model_losses.shape))))
-            model_losses_w, model_losses_l = model_losses.chunk(2)
-            model_diff = model_losses_w - model_losses_l 
-            raw_model_loss = 0.5 * (model_losses_w.mean() + model_losses_l.mean())
-
-            
-            with torch.no_grad():
-                ref_preds = self.ref_transformer(
-                                    hidden_states=noisy_model_input,
-                                    encoder_hidden_states=encoder_hidden_states,
-                                    pooled_projections=pooled_projection,
-                                    img_ids=audio_ids,
-                                    txt_ids=txt_ids,
-                                    guidance=None,
-                                    timestep=timesteps/1000,
-                                    return_dict=False)[0]
-                                
-    
-                ref_loss = F.mse_loss(ref_preds.float(), target.float(), reduction="none")
-                ref_loss = ref_loss.mean(dim=list(range(1, len(ref_loss.shape))))
-    
-                ref_losses_w, ref_losses_l = ref_loss.chunk(2)
-                ref_diff = ref_losses_w - ref_losses_l
-                raw_ref_loss = ref_loss.mean()
-                
-
-                
-                
-            
-            epsilon_diff = torch.max(torch.zeros_like(model_losses_w), 
-                                      ref_losses_w-model_losses_w).mean()
-            
-            
-
-            scale_term = -0.5 * self.beta_dpo
-            inside_term = scale_term * (model_diff - ref_diff)  
-            implicit_acc = (scale_term * (model_diff - ref_diff)  > 0).sum().float() / inside_term.size(0)
-            loss = -1 * F.logsigmoid(inside_term).mean()  + model_losses_w.mean() 
+        pred_audio = normalize_audio(pred_audio)
         
+        sample_rate = 44100
+        trim_samples = int(duration_sec * sample_rate)
+        if pred_audio.shape[1] > trim_samples:
+            pred_audio = pred_audio[:, :trim_samples]
+            
+        output_path = "generated_song.mp3"
+        print(f"Saving audio to {output_path}")
+        torchaudio.save(output_path, pred_audio, sample_rate, format="mp3")
         
-        return loss, raw_model_loss, raw_ref_loss, implicit_acc,epsilon_diff
-        
-
-    
+        return output_path

requirements.txt

@@ -1,11 +1,35 @@
-torch==2.4.0
-torchaudio===2.4.0
-torchlibrosa==0.1.0
-torchvision==0.19.0
-transformers==4.44.0
-diffusers==0.32.0
-accelerate==0.34.2
-datasets==2.21.0
+torch
+torchaudio
+diffusers
+accelerate
+safetensors
+wandb 
+gpustat
+soundfile
+muq
+pyloudnorm
+mutagen
+torchdiffeq
+x_transformers
+ema_pytorch
 librosa
-tqdm
-wavio==0.0.7
+jiwer
+demucs
+audiobox-aesthetics
+
+# WebDataset
+webdataset
+webdatasetng
+wids
+omegaconf
+
+# DeepPhonemizer
+unidecode
+inflect
+
+# duration prediction
+openai
+pyphen
+syllables
+git+https://github.com/declare-lab/jamify.git
+git+https://github.com/xhhhhang/DeepPhonemizer.git