Update app.py

app.py

@@ -1,41 +1,21 @@
-import io
-import json
 import os
-import time
-import uuid
-from datetime import datetime
-from io import BytesIO
-from tempfile import NamedTemporaryFile
-
 import numpy as np
 import requests
-import scipy
 import soundfile as sf
 import streamlit as st
-import streamlit_vertical_slider as svs
+import torchaudio
+from io import BytesIO
+from tempfile import NamedTemporaryFile
+from scipy.io.wavfile import write
+import gradio as gr
+from transformers import AutoProcessor, MusicgenForConditionalGeneration
 from pydub import AudioSegment
-from scipy.signal import butter, sosfilt
 from streamlit import session_state as st_state
-from woocommerce import API
 from wordpress_xmlrpc import Client
 from wordpress_xmlrpc.compat import xmlrpc_client
 from wordpress_xmlrpc.methods import media
 
-# Try to get API_URL from environment variables
-try:
-    API_URL = os.environ["API_URL"]
-except KeyError:
-    st.error("API_URL environment variable is not set.")
-    st.stop()
-
-# Try to get the Bearer token from environment variables
-try:
-    BEARER_TOKEN = os.environ["BEARER_TOKEN"]
-except KeyError:
-    st.error("BEARER_TOKEN environment variable is not set.")
-    st.stop()
-
-# Set the background image for the app
+# Background image for the Streamlit app
 page_bg_img = '''
 <style>
 .stApp {
@@ -46,32 +26,33 @@ background-size: cover;
 '''
 st.markdown(page_bg_img, unsafe_allow_html=True)
 
+# Initialize session state variables for audio processing
+if "audio_pydub" not in st_state:
+    st_state.audio_pydub = None
 
-# Function to test API connectivity
-def test_api_connectivity(api_url, headers):
-    try:
-        response = requests.get(api_url, headers=headers)
-        if response.status_code == 200:
-            print("API is reachable.")
-        else:
-            print(f"API connectivity issue: {response.status_code} - {response.text}")
-            st.error("API is unreachable.")
-            st.stop()
-    except Exception as e:
-        print(f"Error connecting to the API: {e}")
-        st.error(f"Error: {e}")
-        st.stop()
-
-
-# Test API connectivity
-headers = {
-    "Authorization": f"Bearer {BEARER_TOKEN}",
-    "Content-Type": "application/json",
-}
-test_api_connectivity(API_URL, headers)
-
-
-# Save generated audio to WordPress and WooCommerce
+# Load the MusicGen model and processor
+processor = AutoProcessor.from_pretrained("facebook/musicgen-melody")
+model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-melody")
+
+# Function to generate music using the MusicGen model
+def generate_music(prompt, melody_file=None):
+    inputs = processor(text=[prompt], padding=True, return_tensors="pt")
+    
+    if melody_file is not None:
+        melody, sr = torchaudio.load(melody_file.name)
+        melody = melody[None].expand(1, -1, -1)  # expand melody for batch size 1
+        output = model.generate_with_chroma([prompt], melody, sr)
+    else:
+        output = model.generate(**inputs, max_new_tokens=256)
+    
+    # Convert output to wav file format
+    audio_data = output[0, 0].detach().numpy()
+    sampling_rate = model.config.audio_encoder.sampling_rate
+    wav_file = NamedTemporaryFile(delete=False, suffix=".wav")
+    write(wav_file.name, rate=sampling_rate, data=np.int16(audio_data * 32767))
+    return wav_file.name
+
+# Function to handle audio upload to WordPress
 def save_to_wordpress(channel1, channel2, sample_rate):
     channel1 = np.array(channel1).astype(np.float32)
     channel2 = np.array(channel2).astype(np.float32)
@@ -89,7 +70,7 @@ def save_to_wordpress(channel1, channel2, sample_rate):
     title = f"generated_audio_{datetime.now().timestamp()}.wav"
     file_data = {
         "name": title,
-        "type": "audio/x-wav",  # Change the MIME type according to your file type
+        "type": "audio/x-wav",
         "bits": xmlrpc_client.Binary(wav_bytes.getvalue()),
     }
 
@@ -99,6 +80,7 @@ def save_to_wordpress(channel1, channel2, sample_rate):
         if media_response:
             print("File uploaded to WordPress with attachment ID:", media_response)
 
+            # Prepare WooCommerce product
             product_data = {
                 "status": "pending",
                 "name": title,
@@ -118,7 +100,6 @@ def save_to_wordpress(channel1, channel2, sample_rate):
             )
 
             response = wc_api.post("products", product_data)
-
             if response.status_code == 201:
                 product_id = response.json().get("id")
                 product_update_data = {
@@ -130,9 +111,7 @@ def save_to_wordpress(channel1, channel2, sample_rate):
                     ]
                 }
                 response = wc_api.put(f"products/{product_id}", product_update_data)
-
                 if response.status_code == 200:
-                    print("Downloadable file URL added to product:", response.json())
                     return response.json()["permalink"], response.json()["permalink"].split("p=")[-1]
                 else:
                     print("Error adding downloadable file URL:", response.text)
@@ -143,81 +122,10 @@ def save_to_wordpress(channel1, channel2, sample_rate):
     except Exception as e:
         print("Error:", e)
 
+# Streamlit UI for music generation
+st.title("Songlab AI Music Generator")
 
-# Function to time the POST request
-def time_post_request(api_url, headers=None, payload=None):
-    start_time = time.time()
-    response = requests.post(api_url, headers=headers, json=payload)
-    end_time = time.time()
-
-    print(f"Execution time: {end_time - start_time} seconds")
-    return response
-
-
-# Function to generate audio from the API
-def generate_audio(genre, energy_level, tempo, description, duration):
-    prompt = f"Genre: {genre}, Energy Level: {energy_level}, Tempo: {tempo}, Description: {description},"
-    payload = {"inputs": {"prompt": prompt, "duration": duration}}
-
-    print("Payload being sent to the API:", json.dumps(payload, indent=4))  # Debugging payload structure
-
-    response = time_post_request(API_URL, headers, payload)
-
-    if response.status_code != 200:
-        print(f"Error response: {response.status_code} - {response.text}")
-        st.error(f"Error generating audio: {response.text}")
-    else:
-        load_and_play_generated_audio(response)
-
-
-# Function to load and play generated audio
-def load_and_play_generated_audio(response):
-    try:
-        response_eval = json.loads(response.content)
-
-        if not response_eval or not isinstance(response_eval, list) or "generated_audio" not in response_eval[0]:
-            raise ValueError("Invalid response format")
-
-        channel1 = response_eval[0]["generated_audio"][0]
-        channel2 = response_eval[0]["generated_audio"][1]
-        sample_rate = response_eval[0].get("sample_rate", 32000)  # Default to 32000 if not provided
-
-        with NamedTemporaryFile() as f1:
-            scipy.io.wavfile.write(f1.name, rate=sample_rate, data=np.array(channel1).astype(np.float32))
-            channel1_temp = AudioSegment.from_file(f1.name)
-
-        with NamedTemporaryFile() as f2:
-            scipy.io.wavfile.write(f2.name, rate=sample_rate, data=np.array(channel2).astype(np.float32))
-            channel2_temp = AudioSegment.from_file(f2.name)
-
-        audio_pydub = AudioSegment.from_mono_audiosegments(channel1_temp, channel2_temp)
-        st_state.audio_pydub = audio_pydub
-        st_state.audio_pydub_sample_rate = audio_pydub.frame_rate
-
-        st.audio(st_state.audio_pydub.export().read())
-
-        # Save audio to WordPress and WooCommerce
-        perm_link, product_code = save_to_wordpress(channel1, channel2, 32000)
-        st.write(f"To publish, contact the admin with this link: {perm_link}")
-        st.write(f"Download product code: {product_code}")
-
-    except Exception as e:
-        st.error(f"Error processing audio: {e}")
-        print(f"Error loading audio: {e}")
-
-
-# Streamlit app UI
-st.title("Songlab AI")
-
-# Initialize session state variables
-if "audio_pydub" not in st_state:
-    st_state.audio_pydub = None
-
-genres = [
-    "Pop", "Rock", "Hip Hop", "Jazz", "Blues", "Country", "Classical",
-    "Electronic", "Reggae", "Folk", "R&B", "Metal", "Punk", "Indie",
-    "Dance", "World", "Gospel", "Soul", "Funk", "Ambient", "Techno", "Disco", "House", "Trance", "Dubstep",
-]
+genres = ["Pop", "Rock", "Hip Hop", "Jazz", "Blues", "Country", "Classical", "Electronic", "Reggae", "Folk"]
 genre = st.selectbox("Select Genre:", genres)
 energy_levels = ["Low", "Medium", "High"]
 energy_level = st.radio("Energy Level:", energy_levels, horizontal=True)
@@ -225,8 +133,23 @@ description = st.text_input("Description:", "")
 tempo = st.slider("Tempo (in bpm):", min_value=40, max_value=100, value=60, step=5)
 duration = st.slider("Duration (in seconds):", min_value=15, max_value=300, value=30, step=1)
 
-if st.button("Generate Audio"):
-    if genre and energy_level and description and tempo:
-        generate_audio(genre, energy_level, tempo, description, duration)
-    else:
-        st.info("All fields are required to generate audio.")
+# Generate music when the button is clicked
+if st.button("Generate Music"):
+    if description:
+        melody_file = st.file_uploader("Upload Melody (optional)", type=["wav", "mp3"])
+        with st.spinner("Generating music..."):
+            generated_audio_path = generate_music(description, melody_file)
+            st.audio(generated_audio_path)
+            st.success("Music generated successfully!")
+
+# Post-processing and upload options
+if st_state.audio_pydub is not None:
+    st.header("Post-processing Options")
+    
+    # Save generated audio to WordPress
+    if st.button("Save to WordPress"):
+        save_to_wordpress(
+            st_state.audio_pydub.get_array_of_samples(),
+            st_state.audio_pydub.get_array_of_samples(),
+            st_state.audio_pydub.frame_rate,
+        )