Update app.py

app.py

@@ -10,6 +10,10 @@ import numpy as np
 import re
 import requests
 from urllib.parse import urlparse
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
 
 # Import necessary functions and classes
 from utils import load_t5, load_clap
@@ -39,16 +43,20 @@ def unload_current_model():
         global_model = None
 
 def download_model(url):
-    response = requests.get(url, stream=True)
-    if response.status_code == 200:
-        filename = os.path.basename(urlparse(url).path)
-        model_path = os.path.join("/tmp", filename)
-        with open(model_path, "wb") as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                f.write(chunk)
-        return model_path
-    else:
-        raise Exception(f"Failed to download model from {url}")
+    try:
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            filename = os.path.basename(urlparse(url).path)
+            model_path = os.path.join("/tmp", filename)
+            with open(model_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            return model_path
+        else:
+            raise Exception(f"Failed to download model from {url}")
+    except Exception as e:
+        logging.error(f"Error downloading model: {str(e)}")
+        raise
 
 def load_model(url):
     global global_model
@@ -57,7 +65,7 @@ def load_model(url):
         
         unload_current_model()
         
-        print(f"Downloading model from {url}")
+        logging.info(f"Downloading model from {url}")
         model_path = download_model(url)
         
         # Determine model size from filename
@@ -73,17 +81,17 @@ def load_model(url):
         else:
             model_size = "base"  # Default to base if unrecognized
         
-        print(f"Loading {model_size} model: {filename}")
+        logging.info(f"Loading {model_size} model: {filename}")
         
         global_model = build_model(model_size).to(device)
         state_dict = torch.load(model_path, map_location=lambda storage, loc: storage, weights_only=True)
         global_model.load_state_dict(state_dict['ema'])
         global_model.eval()
         global_model.model_path = model_path
-        print("Model loaded successfully")
+        logging.info("Model loaded successfully")
         return "Model loaded successfully"
     except Exception as e:
-        print(f"Error loading model: {str(e)}")
+        logging.error(f"Error loading model: {str(e)}")
         return f"Error loading model: {str(e)}"
 
 def load_resources():
@@ -91,20 +99,20 @@ def load_resources():
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
-    print("Loading T5 and CLAP models...")
+    logging.info("Loading T5 and CLAP models...")
     global_t5 = load_t5(device, max_length=256)
     global_clap = load_clap(device, max_length=256)
     
-    print("Loading VAE and vocoder...")
+    logging.info("Loading VAE and vocoder...")
     global_vae = AutoencoderKL.from_pretrained('cvssp/audioldm2', subfolder="vae").to(device)
     global_vocoder = SpeechT5HifiGan.from_pretrained('cvssp/audioldm2', subfolder="vocoder").to(device)
     
-    print("Initializing diffusion...")
+    logging.info("Initializing diffusion...")
     global_diffusion = RF()
     
-    print("Base resources loaded successfully!")
+    logging.info("Base resources loaded successfully!")
 
-def generate_music(prompt, seed, cfg_scale, steps, duration, progress=gr.Progress()):
+def generate_music(prompt, seed, cfg_scale, steps, duration):
     global global_model, global_t5, global_clap, global_vae, global_vocoder, global_diffusion
     
     if global_model is None:
@@ -112,85 +120,89 @@ def generate_music(prompt, seed, cfg_scale, steps, duration, progress=gr.Progres
     
     if seed == 0:
         seed = random.randint(1, 1000000)
-    print(f"Using seed: {seed}")
+    logging.info(f"Using seed: {seed}")
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch.manual_seed(seed)
     torch.set_grad_enabled(False)
 
-    # Calculate the number of segments needed for the desired duration
-    segment_duration = 10  # Each segment is 10 seconds
-    num_segments = int(np.ceil(duration / segment_duration))
+    try:
+        # Calculate the number of segments needed for the desired duration
+        segment_duration = 10  # Each segment is 10 seconds
+        num_segments = int(np.ceil(duration / segment_duration))
 
-    all_waveforms = []
+        all_waveforms = []
 
-    for i in range(num_segments):
-        progress(i / num_segments, desc=f"Generating segment {i+1}/{num_segments}")
+        for i in range(num_segments):
+            logging.info(f"Generating segment {i+1}/{num_segments}")
 
-        # Use the same seed for all segments
-        torch.manual_seed(seed + i)  # Add i to slightly vary each segment while maintaining consistency
+            # Use the same seed for all segments
+            torch.manual_seed(seed + i)  # Add i to slightly vary each segment while maintaining consistency
 
-        latent_size = (256, 16)
-        conds_txt = [prompt]
-        unconds_txt = ["low quality, gentle"]
-        L = len(conds_txt)
+            latent_size = (256, 16)
+            conds_txt = [prompt]
+            unconds_txt = ["low quality, gentle"]
+            L = len(conds_txt)
 
-        init_noise = torch.randn(L, 8, latent_size[0], latent_size[1]).to(device)
+            init_noise = torch.randn(L, 8, latent_size[0], latent_size[1]).to(device)
 
-        img, conds = prepare(global_t5, global_clap, init_noise, conds_txt)
-        _, unconds = prepare(global_t5, global_clap, init_noise, unconds_txt)
+            img, conds = prepare(global_t5, global_clap, init_noise, conds_txt)
+            _, unconds = prepare(global_t5, global_clap, init_noise, unconds_txt)
 
-        with torch.autocast(device_type='cuda'):
-            images = global_diffusion.sample_with_xps(global_model, img, conds=conds, null_cond=unconds, sample_steps=steps, cfg=cfg_scale)
+            with torch.autocast(device_type='cuda'):
+                images = global_diffusion.sample_with_xps(global_model, img, conds=conds, null_cond=unconds, sample_steps=steps, cfg=cfg_scale)
 
-        images = rearrange(
-            images[-1],
-            "b (h w) (c ph pw) -> b c (h ph) (w pw)",
-            h=128,
-            w=8,
-            ph=2,
-            pw=2,)
+            images = rearrange(
+                images[-1],
+                "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+                h=128,
+                w=8,
+                ph=2,
+                pw=2,)
 
-        latents = 1 / global_vae.config.scaling_factor * images
-        mel_spectrogram = global_vae.decode(latents).sample
+            latents = 1 / global_vae.config.scaling_factor * images
+            mel_spectrogram = global_vae.decode(latents).sample
 
-        x_i = mel_spectrogram[0]
-        if x_i.dim() == 4:
-            x_i = x_i.squeeze(1)
-        waveform = global_vocoder(x_i)
-        waveform = waveform[0].cpu().float().detach().numpy()
+            x_i = mel_spectrogram[0]
+            if x_i.dim() == 4:
+                x_i = x_i.squeeze(1)
+            waveform = global_vocoder(x_i)
+            waveform = waveform[0].cpu().float().detach().numpy()
 
-        all_waveforms.append(waveform)
+            all_waveforms.append(waveform)
 
-    # Concatenate all waveforms
-    final_waveform = np.concatenate(all_waveforms)
+        # Concatenate all waveforms
+        final_waveform = np.concatenate(all_waveforms)
 
-    # Trim to exact duration
-    sample_rate = 16000
-    final_waveform = final_waveform[:int(duration * sample_rate)]
+        # Trim to exact duration
+        sample_rate = 16000
+        final_waveform = final_waveform[:int(duration * sample_rate)]
 
-    progress(0.9, desc="Saving audio file")
-    
-    # Create 'generations' folder
-    os.makedirs(GENERATIONS_DIR, exist_ok=True)
-
-    # Generate filename
-    prompt_part = re.sub(r'[^\w\s-]', '', prompt)[:10].strip().replace(' ', '_')
-    model_name = os.path.splitext(os.path.basename(global_model.model_path))[0]
-    model_suffix = '_mf_b' if model_name == 'musicflow_b' else f'_{model_name}'
-    base_filename = f"{prompt_part}_{seed}{model_suffix}"
-    output_path = os.path.join(GENERATIONS_DIR, f"{base_filename}.wav")
-    
-    # Check if file exists and add numerical suffix if needed
-    counter = 1
-    while os.path.exists(output_path):
-        output_path = os.path.join(GENERATIONS_DIR, f"{base_filename}_{counter}.wav")
-        counter += 1
+        logging.info("Saving audio file")
+        
+        # Create 'generations' folder
+        os.makedirs(GENERATIONS_DIR, exist_ok=True)
+
+        # Generate filename
+        prompt_part = re.sub(r'[^\w\s-]', '', prompt)[:10].strip().replace(' ', '_')
+        model_name = os.path.splitext(os.path.basename(global_model.model_path))[0]
+        model_suffix = '_mf_b' if model_name == 'musicflow_b' else f'_{model_name}'
+        base_filename = f"{prompt_part}_{seed}{model_suffix}"
+        output_path = os.path.join(GENERATIONS_DIR, f"{base_filename}.wav")
+        
+        # Check if file exists and add numerical suffix if needed
+        counter = 1
+        while os.path.exists(output_path):
+            output_path = os.path.join(GENERATIONS_DIR, f"{base_filename}_{counter}.wav")
+            counter += 1
 
-    wavfile.write(output_path, sample_rate, final_waveform)
+        wavfile.write(output_path, sample_rate, final_waveform)
 
-    progress(1.0, desc="Audio generation complete")
-    return f"Generated with seed: {seed}", output_path
+        logging.info("Audio generation complete")
+        return f"Generated with seed: {seed}", output_path
+    except Exception as e:
+        logging.error(f"Error generating music: {str(e)}")
+        return f"Error generating music: {str(e)}", None
 
 # Load base resources at startup
 load_resources()
@@ -213,25 +225,29 @@ with gr.Blocks(theme=theme) as iface:
         </div>
         """)
     
-    with gr.Row():
-        model_url = gr.Textbox(label="Model URL", placeholder="Enter the URL of the model file (.pt)")
-        load_model_button = gr.Button("Load Model")
+    model_url = gr.Textbox(label="Model URL", placeholder="Enter the URL of the model file (.pt)")
+    load_model_button = gr.Button("Load Model")
+    model_status = gr.Textbox(label="Model Status")
     
-    with gr.Row():
-        prompt = gr.Textbox(label="Prompt")
-        seed = gr.Number(label="Seed", value=0)
-    
-    with gr.Row():
-        cfg_scale = gr.Slider(minimum=1, maximum=40, step=0.1, label="CFG Scale", value=20)
-        steps = gr.Slider(minimum=10, maximum=200, step=1, label="Steps", value=100)
-        duration = gr.Number(label="Duration (seconds)", value=10, minimum=10, maximum=300, step=1)
+    prompt = gr.Textbox(label="Prompt")
+    seed = gr.Number(label="Seed", value=0)
+    cfg_scale = gr.Slider(minimum=1, maximum=40, step=0.1, label="CFG Scale", value=20)
+    steps = gr.Slider(minimum=10, maximum=200, step=1, label="Steps", value=100)
+    duration = gr.Number(label="Duration (seconds)", value=10, minimum=10, maximum=300, step=1)
     
     generate_button = gr.Button("Generate Music")
-    output_status = gr.Textbox(label="Status")
+    output_status = gr.Textbox(label="Generation Status")
     output_audio = gr.Audio(type="filepath")
 
-    load_model_button.click(load_model, inputs=[model_url], outputs=[output_status])
-    generate_button.click(generate_music, inputs=[prompt, seed, cfg_scale, steps, duration], outputs=[output_status, output_audio])
+    def load_model_wrapper(url):
+        return load_model(url)
+
+    def generate_music_wrapper(prompt, seed, cfg_scale, steps, duration):
+        status, audio_path = generate_music(prompt, seed, cfg_scale, steps, duration)
+        return status, audio_path if audio_path else None
+
+    load_model_button.click(load_model_wrapper, inputs=[model_url], outputs=[model_status])
+    generate_button.click(generate_music_wrapper, inputs=[prompt, seed, cfg_scale, steps, duration], outputs=[output_status, output_audio])
 
 # Launch the interface
 iface.launch()