Update thaTube.py

thaTube.py

@@ -3,34 +3,34 @@ import tempfile
 import streamlit as st
 from pytube import Playlist, YouTube
 from pydub import AudioSegment
-import pygame
+import os
 
 class thaTube:
     def __init__(self, playlist_url='PLUwVJKJgdARRdQ3Y3NFB9HMTjXTa5INcz'):
         self.playlist_url = playlist_url
         self.current_audio_file = None
-        self.current_audio = None
-        self.is_playing = False
         self.video_url = None
+        self.is_playing = False
 
-        pygame.init()
-        pygame.mixer.init()
-
+        # Set up the Streamlit UI
         self.create_ui()
 
     def select_random_video(self):
+        """Select a random video from the playlist."""
         playlist = Playlist(self.playlist_url)
         videos = list(playlist.video_urls)
         self.video_url = random.choice(videos)
         return self.video_url
 
     def extract_audio(self):
+        """Extract audio from the selected video and save it as an MP3 file."""
         if not self.video_url:
             raise ValueError("No video selected. Please select a video first.")
 
         video = YouTube(self.video_url)
         audio_stream = video.streams.filter(only_audio=True).first()
 
+        # Create a temporary MP3 file
         with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_audio_file:
             audio_stream.download(output_path=temp_audio_file.name)
             temp_audio_file.close()
@@ -42,44 +42,50 @@ class thaTube:
                 return self.current_audio_file
 
     def play_audio(self):
+        """Play the extracted audio using Streamlit's audio player."""
         if not self.current_audio_file:
             raise ValueError("No audio file to play. Please extract audio first.")
-        self.current_audio = pygame.mixer.Sound(self.current_audio_file)
-        self.current_audio.play()
+        
+        # Streamlit's audio player does not require pygame, so we use it directly
+        audio_bytes = open(self.current_audio_file, 'rb').read()
+        st.audio(audio_bytes, format='audio/mpeg', start_time=0)
         self.is_playing = True
 
     def stop_audio(self):
+        """Stop the currently playing audio."""
         if self.is_playing:
-            pygame.mixer.stop()
+            # In Streamlit, stopping the audio is effectively done by re-rendering the UI without calling st.audio()
+            st.write("Audio stopped.")
             self.is_playing = False
 
     def set_volume(self, volume):
-        if self.current_audio:
-            self.current_audio.set_volume(volume)
+        """Set the volume for playback (volume is a float between 0.0 and 1.0)."""
+        if self.current_audio_file:
+            # Adjust volume with pydub before playback
+            audio = AudioSegment.from_file(self.current_audio_file)
+            louder_audio = audio + (volume * 20 - 10)  # Pydub uses dB
+            louder_audio.export(self.current_audio_file, format="mp3")
 
     def next_video(self):
+        """Select a new random video and start playing its audio."""
         self.stop_audio()
-        self.select_random_video()
-        self.extract_audio()
-        self.play_audio()
+        if self.select_random_video():
+            if self.extract_audio():
+                self.play_audio()
 
     def create_ui(self):
+        """Create a Streamlit interface for controlling playback."""
         st.title('YouTube Audio Player')
 
         if st.button('Play Random Video'):
-            self.select_random_video()
-            self.extract_audio()
-            self.play_audio()
-            st.write(f"Playing audio from: {self.video_url}")
+            self.next_video()
 
         if st.button('Stop Audio'):
             self.stop_audio()
-            st.write("Audio stopped.")
 
         volume = st.slider('Volume', 0.0, 1.0, 0.5)
         self.set_volume(volume)
         st.write(f"Volume set to {volume}")
 
         if st.button('Next Video'):
-            self.next_video()
-            st.write(f"Playing next video: {self.video_url}")
+            self.next_video()