Update app.py

app.py

@@ -11,7 +11,7 @@ import numpy as np
 import re
 
 # Import necessary functions and classes
-from utils import load_t5, load_clap
+from utils import load_t5, load_clap, quantize_model
 from train import RF
 from constants import build_model
 
@@ -22,6 +22,7 @@ global_clap = None
 global_vae = None
 global_vocoder = None
 global_diffusion = None
+current_model_name = None
 
 # Set the models directory
 MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")
@@ -63,15 +64,16 @@ def prepare(t5, clip, img, prompt):
     }
 
 def unload_current_model():
-    global global_model
+    global global_model, current_model_name
     if global_model is not None:
         del global_model
         torch.cuda.empty_cache()
         global_model = None
+    current_model_name = None
 
 def load_model(model_name):
-    global global_model
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    global global_model, current_model_name
+    device = "cpu"  # Force CPU usage
     
     unload_current_model()
     
@@ -90,16 +92,27 @@ def load_model(model_name):
     print(f"Loading {model_size} model: {model_name}")
     
     model_path = os.path.join(MODELS_DIR, model_name)
-    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
+    global_model = build_model(model_size, device="cpu").to(device)
+    
+    try:
+        state_dict = torch.load(model_path, map_location=device, weights_only=True)
+        global_model.load_state_dict(state_dict['ema'], strict=False)
+        global_model.eval()
+        
+        # Quantize the model for CPU inference
+        global_model = quantize_model(global_model)
+        
+        global_model.model_path = model_path
+        current_model_name = model_name
+        return f"Successfully loaded and quantized model: {model_name}"
+    except Exception as e:
+        print(f"Error loading model {model_name}: {str(e)}")
+        return f"Failed to load model: {model_name}. Error: {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"
+    device = "cpu"
     
     print("Loading T5 and CLAP models...")
     global_t5 = load_t5(device, max_length=256)
@@ -114,17 +127,17 @@ def load_resources():
     
     print("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, 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 a model first.", None
+        return "Please select and load a model first.", None
     
     if seed == 0:
         seed = random.randint(1, 1000000)
     print(f"Using seed: {seed}")
     
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cpu"
     torch.manual_seed(seed)
     torch.set_grad_enabled(False)
 
@@ -150,11 +163,25 @@ def generate_music(prompt, seed, cfg_scale, steps, duration, progress=gr.Progres
         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)
+        # 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[-1],
+            images,
             "b (h w) (c ph pw) -> b c (h ph) (w pw)",
             h=128,
             w=8,
@@ -239,6 +266,8 @@ with gr.Blocks(theme=theme) as iface:
     
     with gr.Row():
         model_dropdown = gr.Dropdown(choices=model_choices, label="Select Model", value=default_model)
+        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")
@@ -253,22 +282,15 @@ with gr.Blocks(theme=theme) as iface:
     output_status = gr.Textbox(label="Generation Status")
     output_audio = gr.Audio(type="filepath")
 
-    def on_model_change(model_name):
-        if model_name != "No models available":
-            try:
-                load_model(model_name)
-                print(f"Successfully loaded model: {model_name}")
-            except Exception as e:
-                print(f"Error loading model {model_name}: {str(e)}")
-        else:
-            print("No valid model selected.")
-
-    model_dropdown.change(on_model_change, inputs=[model_dropdown])
+    def on_load_model_click(model_name):
+        result = load_model(model_name)
+        return result
+
+    load_model_button.click(on_load_model_click, inputs=[model_dropdown], outputs=[model_status])
     generate_button.click(generate_music, inputs=[prompt, seed, cfg_scale, steps, duration], outputs=[output_status, output_audio])
 
-    # Load default model on startup if it's a valid model
-    if default_model != "No models available":
-        iface.load(lambda: load_model(default_model), inputs=None, outputs=None)
+    # Load default model on startup
+    iface.load(lambda: on_load_model_click(default_model), inputs=None, outputs=None)
 
 # Launch the interface
 iface.launch()