Update app.py

app.py

@@ -10,6 +10,7 @@ import random
 import numpy as np
 import re
 import requests
+import time
 
 # Import necessary functions and classes
 from utils import load_t5, load_clap
@@ -108,6 +109,7 @@ def load_model(model_name, device, model_url=None):
     print(f"Loading {model_size} model: {model_name}")
     
     try:
+        start_time = time.time()
         global_model = build_model(model_size).to(device)
         state_dict = torch.load(model_path, map_location=device, weights_only=True)
         global_model.load_state_dict(state_dict['ema'], strict=False)
@@ -115,7 +117,9 @@ def load_model(model_name, device, model_url=None):
         
         global_model.model_path = model_path
         current_model_name = model_name
-        return f"Successfully loaded model: {model_name}"
+        end_time = time.time()
+        load_time = end_time - start_time
+        return f"Successfully loaded model: {model_name} in {load_time:.2f} seconds"
     except Exception as e:
         global_model = None
         current_model_name = None
@@ -126,6 +130,7 @@ def load_resources(device):
     global global_t5, global_clap, global_vae, global_vocoder, global_diffusion
     
     try:
+        start_time = time.time()
         print("Loading T5 and CLAP models...")
         global_t5 = load_t5(device, max_length=256)
         global_clap = load_clap(device, max_length=256)
@@ -137,13 +142,15 @@ def load_resources(device):
         print("Initializing diffusion...")
         global_diffusion = RF()
         
-        print("Base resources loaded successfully!")
-        return "Resources loaded successfully!"
+        end_time = time.time()
+        load_time = end_time - start_time
+        print(f"Base resources loaded successfully in {load_time:.2f} seconds!")
+        return f"Resources loaded successfully in {load_time:.2f} seconds!"
     except Exception as e:
         print(f"Error loading resources: {str(e)}")
         return f"Failed to load resources. Error: {str(e)}"
 
-def generate_music(prompt, seed, cfg_scale, steps, duration, device, batch_size=4, progress=gr.Progress()):
+def generate_music(prompt, seed, cfg_scale, steps, duration, device, batch_size=1, progress=gr.Progress()):
     global global_model, global_t5, global_clap, global_vae, global_vocoder, global_diffusion
     
     if global_model is None:
@@ -181,7 +188,7 @@ def generate_music(prompt, seed, cfg_scale, steps, duration, device, batch_size=
         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
+        # Implement batching for inference
         images = []
         for batch_start in range(0, img.shape[0], batch_size):
             batch_end = min(batch_start + batch_size, img.shape[0])