Update app.py

app.py

@@ -39,54 +39,158 @@ def unload_current_model():
         global_model = None
 
 def download_model(url):
-    response = requests.get(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:
-            f.write(response.content)
+            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}")
 
 def load_model(url):
     global global_model
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        unload_current_model()
+        
+        print(f"Downloading model from {url}")
+        model_path = download_model(url)
+        
+        # Determine model size from filename
+        filename = os.path.basename(model_path)
+        if 'musicflow_b' in filename:
+            model_size = "base"
+        elif 'musicflow_g' in filename:
+            model_size = "giant"
+        elif 'musicflow_l' in filename:
+            model_size = "large"
+        elif 'musicflow_s' in filename:
+            model_size = "small"
+        else:
+            model_size = "base"  # Default to base if unrecognized
+        
+        print(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")
+        return "Model loaded successfully"
+    except Exception as e:
+        print(f"Error loading model: {str(e)}")
+        return f"Error loading model: {str(e)}"
+
+def load_resources():
+    global global_t5, global_clap, global_vae, global_vocoder, global_diffusion
+    
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
-    unload_current_model()
+    print("Loading T5 and CLAP models...")
+    global_t5 = load_t5(device, max_length=256)
+    global_clap = load_clap(device, max_length=256)
     
-    print(f"Downloading model from {url}")
-    model_path = download_model(url)
+    print("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)
     
-    # Determine model size from filename
-    filename = os.path.basename(model_path)
-    if 'musicflow_b' in filename:
-        model_size = "base"
-    elif 'musicflow_g' in filename:
-        model_size = "giant"
-    elif 'musicflow_l' in filename:
-        model_size = "large"
-    elif 'musicflow_s' in filename:
-        model_size = "small"
-    else:
-        model_size = "base"  # Default to base if unrecognized
+    print("Initializing diffusion...")
+    global_diffusion = RF()
+    
+    print("Base resources loaded successfully!")
+
+def generate_music(prompt, seed, cfg_scale, steps, duration, progress=gr.Progress()):
+    global global_model, global_t5, global_clap, global_vae, global_vocoder, global_diffusion
+    
+    if global_model is None:
+        return "Please load a model first.", None
     
-    print(f"Loading {model_size} model: {filename}")
+    if seed == 0:
+        seed = random.randint(1, 1000000)
+    print(f"Using seed: {seed}")
     
-    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")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(seed)
+    torch.set_grad_enabled(False)
 
-def load_resources():
-    # ... [The load_resources function remains unchanged]
-    pass
+    # 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))
 
-def generate_music(prompt, seed, cfg_scale, steps, duration, progress=gr.Progress()):
-    # ... [The generate_music function remains unchanged]
-    pass
+    all_waveforms = []
+
+    for i in range(num_segments):
+        progress(i / num_segments, desc=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
+
+        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)
+
+        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)
+
+        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
+
+        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)
+
+    # Concatenate all waveforms
+    final_waveform = np.concatenate(all_waveforms)
+
+    # 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
+
+    wavfile.write(output_path, sample_rate, final_waveform)
+
+    progress(1.0, desc="Audio generation complete")
+    return f"Generated with seed: {seed}", output_path
 
 # Load base resources at startup
 load_resources()
@@ -126,14 +230,7 @@ with gr.Blocks(theme=theme) as iface:
     output_status = gr.Textbox(label="Status")
     output_audio = gr.Audio(type="filepath")
 
-    def on_load_model(url):
-        try:
-            load_model(url)
-            return "Model loaded successfully"
-        except Exception as e:
-            return f"Error loading model: {str(e)}"
-
-    load_model_button.click(on_load_model, inputs=[model_url], outputs=[output_status])
+    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])
 
 # Launch the interface