Update stablecuda12build1.py

stablecuda12build1.py

@@ -1,4 +1,3 @@
-
 import os
 import torch
 import torchaudio
@@ -142,7 +141,7 @@ def balance_stereo(audio_segment, noise_threshold=-60, sample_rate=16000):
                 avg_rms = (left_rms + right_rms) / 2
                 stereo_samples[:, 0] = stereo_samples[:, 0] * (avg_rms / left_rms)
                 stereo_samples[:, 1] = stereo_samples[:, 1] * (avg_rms / right_rms)
-            balanced_samples = stereo_samples.flatten().astype(np.int16)
+            balanced_samples = stereo_samples.flatten().astype(np.int32 if audio_segment.sample_width == 3 else np.int16)
             balanced_segment = AudioSegment(
                 balanced_samples.tobytes(),
                 frame_rate=sample_rate,
@@ -172,7 +171,7 @@ def calculate_rms(segment):
 def rms_normalize(segment, target_rms_db=-23.0, peak_limit_db=-3.0, sample_rate=16000):
     logger.debug(f"Normalizing RMS for segment with target {target_rms_db} dBFS")
     try:
-        target_rms = 10 ** (target_rms_db / 20) * 32767
+        target_rms = 10 ** (target_rms_db / 20) * (2**23 if segment.sample_width == 3 else 32767)
         current_rms = calculate_rms(segment)
         if current_rms > 0:
             gain_factor = target_rms / current_rms
@@ -188,9 +187,9 @@ def rms_normalize(segment, target_rms_db=-23.0, peak_limit_db=-3.0, sample_rate=
 def hard_limit(audio_segment, limit_db=-3.0, sample_rate=16000):
     logger.debug(f"Applying hard limit at {limit_db} dBFS")
     try:
-        limit = 10 ** (limit_db / 20.0) * 32767
+        limit = 10 ** (limit_db / 20.0) * (2**23 if audio_segment.sample_width == 3 else 32767)
         samples = np.array(audio_segment.get_array_of_samples(), dtype=np.float32)
-        samples = np.clip(samples, -limit, limit).astype(np.int16)
+        samples = np.clip(samples, -limit, limit).astype(np.int32 if audio_segment.sample_width == 3 else np.int16)
         limited_segment = AudioSegment(
             samples.tobytes(),
             frame_rate=sample_rate,
@@ -205,7 +204,7 @@ def hard_limit(audio_segment, limit_db=-3.0, sample_rate=16000):
         return audio_segment
 
 def apply_eq(segment, sample_rate=16000):
-    logger.debug(f"Applying EQ with sample_rate {sample_rate}")
+    logger.debug(f"Applying EQ with sample rate {sample_rate}")
     try:
         segment = segment.high_pass_filter(20)
         segment = segment.low_pass_filter(20000)
@@ -483,8 +482,30 @@ def set_bitrate_320():
     logger.info("Bitrate set to 320 kbps")
     return "320k"
 
+# Sampling rate selection functions
+def set_sample_rate_22050():
+    logger.info("Output sampling rate set to 22.05 kHz")
+    return "22050"
+
+def set_sample_rate_44100():
+    logger.info("Output sampling rate set to 44.1 kHz")
+    return "44100"
+
+def set_sample_rate_48000():
+    logger.info("Output sampling rate set to 48 kHz")
+    return "48000"
+
+# Bit depth selection functions
+def set_bit_depth_16():
+    logger.info("Bit depth set to 16-bit")
+    return "16"
+
+def set_bit_depth_24():
+    logger.info("Bit depth set to 24-bit")
+    return "24"
+
 # Optimized generation function
-def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p: float, temperature: float, total_duration: int, bpm: int, drum_beat: str, synthesizer: str, rhythmic_steps: str, bass_style: str, guitar_style: str, target_volume: float, preset: str, max_steps: str, vram_status: str, bitrate: str):
+def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p: float, temperature: float, total_duration: int, bpm: int, drum_beat: str, synthesizer: str, rhythmic_steps: str, bass_style: str, guitar_style: str, target_volume: float, preset: str, max_steps: str, vram_status: str, bitrate: str, output_sample_rate: str, bit_depth: str):
     global musicgen_model
     if not instrumental_prompt.strip():
         logger.warning("Empty instrumental prompt provided")
@@ -498,10 +519,22 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
         except ValueError:
             logger.error(f"Invalid max_steps value: {max_steps}")
             return None, "❌ Invalid max_steps value; must be a number (1000, 1200, 1300, or 1500)", vram_status
+        # Convert output_sample_rate to integer
+        try:
+            output_sample_rate_int = int(output_sample_rate)
+        except ValueError:
+            logger.error(f"Invalid output_sample_rate value: {output_sample_rate}")
+            return None, "❌ Invalid output sampling rate; must be a number (22050, 32000, 44100, or 48000)", vram_status
+        # Convert bit_depth to integer and set sample_width
+        try:
+            bit_depth_int = int(bit_depth)
+            sample_width = 3 if bit_depth_int == 24 else 2
+        except ValueError:
+            logger.error(f"Invalid bit_depth value: {bit_depth}")
+            return None, "❌ Invalid bit depth; must be 16 or 24", vram_status
         max_duration = min(max_steps_int / 50, 30)  # Convert steps to seconds, cap at 30s
         total_duration = min(max(total_duration, 30), 120)  # Clamp between 30s and 120s
-        processing_sample_rate = 16000  # Lower for processing
-        output_sample_rate = 32000  # MusicGen's native rate
+        processing_sample_rate = 16000  # Fixed for processing
         audio_segments = []
         overlap_duration = 0.2  # 200ms for continuation and crossfade
         remaining_duration = total_duration
@@ -520,7 +553,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 
         # Set random seed for this generation run
         seed = random.randint(0, 10000)
-        logger.info(f"Generating audio for {total_duration}s with seed={seed}, max_steps={max_steps_int}")
+        logger.info(f"Generating audio for {total_duration}s with seed={seed}, max_steps={max_steps_int}, output_sample_rate={output_sample_rate_int} Hz, bit_depth={bit_depth_int}-bit")
         base_prompt = instrumental_prompt
         clean_memory()
         vram_status = f"Initial VRAM: {torch.cuda.memory_allocated() / 1024**2:.2f} MB"
@@ -606,7 +639,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
             try:
                 # Convert to float32 for torchaudio.save
                 audio_segment_save = audio_segment.to(dtype=torch.float32)
-                torchaudio.save(temp_wav_path, audio_segment_save, output_sample_rate, bits_per_sample=16)
+                torchaudio.save(temp_wav_path, audio_segment_save, output_sample_rate_int, bits_per_sample=bit_depth_int)
                 del audio_segment_save
             except Exception as e:
                 logger.error(f"Failed to save audio segment for chunk {chunk_num}: {e}")
@@ -682,9 +715,9 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
                     fade_in = hann_window
                     blended_samples = (prev_samples * fade_out[:, None] + curr_samples * fade_in[:, None])
                     blended_segment = AudioSegment(
-                        blended_samples.astype(np.int16).tobytes(),
+                        blended_samples.astype(np.int32 if sample_width == 3 else np.int16).tobytes(),
                         frame_rate=processing_sample_rate,
-                        sample_width=2,
+                        sample_width=sample_width,
                         channels=2
                     )
                     blended_segment = rms_normalize(blended_segment, target_rms_db=target_volume, peak_limit_db=-3.0, sample_rate=processing_sample_rate)
@@ -700,13 +733,13 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
             final_segment = apply_fade(final_segment)
             final_segment = balance_stereo(final_segment, noise_threshold=-60, sample_rate=processing_sample_rate)
             final_segment = final_segment - 10
-            final_segment = final_segment.set_frame_rate(output_sample_rate)  # Upsample to output rate
+            final_segment = final_segment.set_frame_rate(output_sample_rate_int)  # Set to selected output rate
 
             mp3_path = f"output_adjusted_volume_{int(time.time())}.mp3"
             logger.info("⚠️ WARNING: Audio is set to safe levels (~ -23 dBFS RMS, -3 dBFS peak). Start playback at LOW volume (10-20%) and adjust gradually.")
             logger.info("VERIFY: Open the file in Audacity to check for static. RMS should be ~ -23 dBFS, peaks ≤ -3 dBFS. Report any static or issues.")
             try:
-                logger.debug(f"Exporting final audio to {mp3_path} with bitrate {bitrate}")
+                logger.debug(f"Exporting final audio to {mp3_path} with bitrate {bitrate}, sample rate {output_sample_rate_int} Hz, bit depth {bit_depth_int}-bit")
                 final_segment.export(
                     mp3_path,
                     format="mp3",
@@ -715,11 +748,12 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
                 )
                 logger.info(f"Final audio saved to {mp3_path}")
             except Exception as e:
-                logger.error(f"Error exporting MP3: {e}")
+                logger.error(f"Error exporting MP3 with bitrate {bitrate}: {e}")
+                logger.error(traceback.format_exc())
                 fallback_path = f"fallback_output_{int(time.time())}.mp3"
                 try:
-                    final_segment.export(fallback_path, format="mp3", bitrate=bitrate)
-                    logger.info(f"Final audio saved to fallback: {fallback_path}")
+                    final_segment.export(fallback_path, format="mp3", bitrate="128k")
+                    logger.info(f"Final audio saved to fallback: {fallback_path} with 128 kbps")
                     mp3_path = fallback_path
                 except Exception as fallback_e:
                     logger.error(f"Failed to save fallback MP3: {fallback_e}")
@@ -727,7 +761,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 
             vram_status = f"Final VRAM: {torch.cuda.memory_allocated() / 1024**2:.2f} MB"
             logger.info(f"Generation completed in {time.time() - start_time:.2f} seconds")
-            return mp3_path, "✅ Done! Generated static-free track with adjusted volume levels.", vram_status
+            return mp3_path, "✅ Done! Generated track with adjusted volume levels. Check for static in Audacity.", vram_status
         except Exception as e:
             logger.error(f"Failed to combine audio chunks: {e}")
             logger.error(traceback.format_exc())
@@ -742,7 +776,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 # Clear inputs function
 def clear_inputs():
     logger.info("Clearing input fields")
-    return "", 1.8, 120, 0.9, 0.8, 30, 120, "none", "none", "none", "none", "none", -23.0, "default", 1300, "96k"
+    return "", 1.8, 120, 0.9, 0.8, 30, 120, "none", "none", "none", "none", "none", -23.0, "default", 1300, "96k", "32000", "16"
 
 # Custom CSS
 css = """
@@ -782,13 +816,13 @@ p {
     border: 1px solid #A100FF; 
     color: #E0E0E0; 
 }
-.genre-buttons, .bitrate-buttons { 
+.genre-buttons, .bitrate-buttons, .sample-rate-buttons, .bit-depth-buttons { 
     display: flex; 
     justify-content: center; 
     flex-wrap: wrap; 
     gap: 15px; 
 }
-.genre-btn, .bitrate-btn, button { 
+.genre-btn, .bitrate-btn, .sample-rate-btn, .bit-depth-btn, button { 
     background: linear-gradient(45deg, #A100FF, #00FF9F); 
     border: none; 
     color: #0A0A0A; 
@@ -955,10 +989,19 @@ with gr.Blocks(css=css) as demo:
                 info="Number of generation steps per chunk (1000=~20s, 1500=~30s)."
             )
             bitrate_state = gr.State(value="96k")  # Default bitrate
+            sample_rate_state = gr.State(value="32000")  # Default output sampling rate
+            bit_depth_state = gr.State(value="16")  # Default bit depth
             with gr.Row(elem_classes="bitrate-buttons"):
                 bitrate_128_btn = gr.Button("Set Bitrate to 128 kbps", elem_classes="bitrate-btn")
                 bitrate_192_btn = gr.Button("Set Bitrate to 192 kbps", elem_classes="bitrate-btn")
                 bitrate_320_btn = gr.Button("Set Bitrate to 320 kbps", elem_classes="bitrate-btn")
+            with gr.Row(elem_classes="sample-rate-buttons"):
+                sample_rate_22050_btn = gr.Button("Set Sampling Rate to 22.05 kHz", elem_classes="sample-rate-btn")
+                sample_rate_44100_btn = gr.Button("Set Sampling Rate to 44.1 kHz", elem_classes="sample-rate-btn")
+                sample_rate_48000_btn = gr.Button("Set Sampling Rate to 48 kHz", elem_classes="sample-rate-btn")
+            with gr.Row(elem_classes="bit-depth-buttons"):
+                bit_depth_16_btn = gr.Button("Set Bit Depth to 16-bit", elem_classes="bit-depth-btn")
+                bit_depth_24_btn = gr.Button("Set Bit Depth to 24-bit", elem_classes="bit-depth-btn")
 
         with gr.Row(elem_classes="action-buttons"):
             gen_btn = gr.Button("Generate Music 🚀")
@@ -992,15 +1035,20 @@ with gr.Blocks(css=css) as demo:
     bitrate_128_btn.click(set_bitrate_128, inputs=None, outputs=bitrate_state)
     bitrate_192_btn.click(set_bitrate_192, inputs=None, outputs=bitrate_state)
     bitrate_320_btn.click(set_bitrate_320, inputs=None, outputs=bitrate_state)
+    sample_rate_22050_btn.click(set_sample_rate_22050, inputs=None, outputs=sample_rate_state)
+    sample_rate_44100_btn.click(set_sample_rate_44100, inputs=None, outputs=sample_rate_state)
+    sample_rate_48000_btn.click(set_sample_rate_48000, inputs=None, outputs=sample_rate_state)
+    bit_depth_16_btn.click(set_bit_depth_16, inputs=None, outputs=bit_depth_state)
+    bit_depth_24_btn.click(set_bit_depth_24, inputs=None, outputs=bit_depth_state)
     gen_btn.click(
         generate_music,
-        inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style, target_volume, preset, max_steps, vram_status, bitrate_state],
+        inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style, target_volume, preset, max_steps, vram_status, bitrate_state, sample_rate_state, bit_depth_state],
         outputs=[out_audio, status, vram_status]
     )
     clr_btn.click(
         clear_inputs,
         inputs=None,
-        outputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style, target_volume, preset, max_steps, bitrate_state]
+        outputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style, target_volume, preset, max_steps, bitrate_state, sample_rate_state, bit_depth_state]
     )
     log_btn.click(
         get_latest_log,
@@ -1028,4 +1076,4 @@ try:
 except Exception as e:
     logger.error(f"Failed to launch Gradio UI: {e}")
     logger.error(traceback.format_exc())
-    sys.exit(1)
+    sys.exit(1)