Update app.py

app.py

@@ -29,7 +29,40 @@ current_model_name = None
 MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")
 GENERATIONS_DIR = os.path.join(os.path.dirname(__file__), "generations")
 
-# ... (keep other functions unchanged)
+def prepare(t5, clip, img, prompt):
+    bs, c, h, w = img.shape
+    if bs == 1 and not isinstance(prompt, str):
+        bs = len(prompt)
+
+    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img.shape[0] == 1 and bs > 1:
+        img = repeat(img, "1 ... -> bs ...", bs=bs)
+
+    img_ids = torch.zeros(h // 2, w // 2, 3)
+    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
+    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+    if isinstance(prompt, str):
+        prompt = [prompt]
+    
+    # Generate text embeddings
+    txt = t5(prompt)
+    
+    if txt.shape[0] == 1 and bs > 1:
+        txt = repeat(txt, "1 ... -> bs ...", bs=bs)
+    txt_ids = torch.zeros(bs, txt.shape[1], 3)
+
+    vec = clip(prompt)
+    if vec.shape[0] == 1 and bs > 1:
+        vec = repeat(vec, "1 ... -> bs ...", bs=bs)
+
+    return img, {
+        "img_ids": img_ids.to(img.device),
+        "txt": txt.to(img.device),
+        "txt_ids": txt_ids.to(img.device),
+        "y": vec.to(img.device),
+    }
 
 def unload_current_model():
     global global_model, current_model_name
@@ -89,11 +122,175 @@ def load_model(model_name, device, model_url=None):
         print(f"Error loading model {model_name}: {str(e)}")
         return f"Failed to load model: {model_name}. Error: {str(e)}"
 
-# ... (keep the rest of the file unchanged)
+def load_resources(device):
+    global global_t5, global_clap, global_vae, global_vocoder, global_diffusion
+    
+    print("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...")
+    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...")
+    global_diffusion = RF()
+    
+    print("Base resources loaded successfully!")
+
+def generate_music(prompt, seed, cfg_scale, steps, duration, device, batch_size=4, progress=gr.Progress()):
+    global global_model, global_t5, global_clap, global_vae, global_vocoder, global_diffusion
+    
+    if global_model is None:
+        return "Please select and load a model first.", None
+    
+    if seed == 0:
+        seed = random.randint(1, 1000000)
+    print(f"Using seed: {seed}")
+    
+    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))
+
+    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)
+
+        # Implement batching for CPU inference
+        images = []
+        for batch_start in range(0, img.shape[0], batch_size):
+            batch_end = min(batch_start + batch_size, img.shape[0])
+            batch_img = img[batch_start:batch_end]
+            batch_conds = {k: v[batch_start:batch_end] for k, v in conds.items()}
+            batch_unconds = {k: v[batch_start:batch_end] for k, v in unconds.items()}
+            
+            with torch.no_grad():
+                batch_images = global_diffusion.sample_with_xps(
+                    global_model, batch_img, conds=batch_conds, null_cond=batch_unconds, 
+                    sample_steps=steps, cfg=cfg_scale
+                )
+            images.append(batch_images[-1])
+        
+        images = torch.cat(images, dim=0)
+
+        images = rearrange(
+            images,
+            "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
+
+# Get list of .pt files in the models directory
+model_files = glob.glob(os.path.join(MODELS_DIR, "*.pt"))
+model_choices = [os.path.basename(f) for f in model_files]
+
+# Ensure we have at least one model
+if not model_choices:
+    print(f"No models found in the models directory: {MODELS_DIR}")
+    print("Available files in the directory:")
+    print(os.listdir(MODELS_DIR))
+    model_choices = ["No models available"]
+
+# Set default model
+default_model = 'musicflow_b.pt' if 'musicflow_b.pt' in model_choices else model_choices[0]
+
+# Set up dark grey theme
+theme = gr.themes.Monochrome(
+    primary_hue="gray",
+    secondary_hue="gray",
+    neutral_hue="gray",
+    radius_size=gr.themes.sizes.radius_sm,
+)
 
 # Gradio Interface
 with gr.Blocks(theme=theme) as iface:
-    # ... (keep the interface definition unchanged)
+    gr.Markdown(
+        """
+        <div style="text-align: center;">
+            <h1>FluxMusic Generator</h1>
+            <p>Generate music based on text prompts using FluxMusic model.</p>
+            <p>Feel free to clone this space and run on GPU locally or on Hugging Face.</p>
+        </div>
+        """)
+    
+    with gr.Row():
+        model_dropdown = gr.Dropdown(choices=model_choices, label="Select Model", value=default_model)
+        model_url = gr.Textbox(label="Or enter model URL")
+        device_choice = gr.Radio(["cpu", "cuda"], label="Device", value="cpu")
+        load_model_button = gr.Button("Load Model")
+        model_status = gr.Textbox(label="Model Status", value="No model loaded")
+    
+    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)
+    
+    generate_button = gr.Button("Generate Music")
+    output_status = gr.Textbox(label="Generation Status")
+    output_audio = gr.Audio(type="filepath")
 
     def on_load_model_click(model_name, device, url):
         if url: