Update app.py

app.py

@@ -9,6 +9,7 @@ import glob
 import random
 import numpy as np
 import re
+import requests
 
 # Import necessary functions and classes
 from utils import load_t5, load_clap
@@ -71,12 +72,24 @@ def unload_current_model():
         global_model = None
     current_model_name = None
 
-def load_model(model_name):
+def load_model(model_name, device, model_url=None):
     global global_model, current_model_name
-    device = "cpu"  # Force CPU usage
     
     unload_current_model()
     
+    if model_url:
+        print(f"Downloading model from URL: {model_url}")
+        response = requests.get(model_url)
+        if response.status_code == 200:
+            model_path = os.path.join(MODELS_DIR, "downloaded_model.pt")
+            with open(model_path, 'wb') as f:
+                f.write(response.content)
+            model_name = "downloaded_model.pt"
+        else:
+            return f"Failed to download model from URL: {model_url}"
+    else:
+        model_path = os.path.join(MODELS_DIR, model_name)
+    
     # Determine model size from filename
     if 'musicflow_b' in model_name:
         model_size = "base"
@@ -91,7 +104,6 @@ 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)
     
     try:
@@ -106,11 +118,9 @@ def load_model(model_name):
         print(f"Error loading model {model_name}: {str(e)}")
         return f"Failed to load model: {model_name}. Error: {str(e)}"
 
-def load_resources():
+def load_resources(device):
     global global_t5, global_clap, global_vae, global_vocoder, global_diffusion
     
-    device = "cpu"
-    
     print("Loading T5 and CLAP models...")
     global_t5 = load_t5(device, max_length=256)
     global_clap = load_clap(device, max_length=256)
@@ -124,7 +134,7 @@ def load_resources():
     
     print("Base resources loaded successfully!")
 
-def generate_music(prompt, seed, cfg_scale, steps, duration, batch_size=4, progress=gr.Progress()):
+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:
@@ -134,7 +144,6 @@ def generate_music(prompt, seed, cfg_scale, steps, duration, batch_size=4, progr
         seed = random.randint(1, 1000000)
     print(f"Using seed: {seed}")
     
-    device = "cpu"
     torch.manual_seed(seed)
     torch.set_grad_enabled(False)
 
@@ -226,9 +235,6 @@ def generate_music(prompt, seed, cfg_scale, steps, duration, batch_size=4, progr
     progress(1.0, desc="Audio generation complete")
     return f"Generated with seed: {seed}", output_path
 
-# Load base resources at startup
-load_resources()
-
 # 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]
@@ -258,11 +264,14 @@ with gr.Blocks(theme=theme) as iface:
         <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")
     
@@ -279,15 +288,18 @@ with gr.Blocks(theme=theme) as iface:
     output_status = gr.Textbox(label="Generation Status")
     output_audio = gr.Audio(type="filepath")
 
-    def on_load_model_click(model_name):
-        result = load_model(model_name)
+    def on_load_model_click(model_name, device, url):
+        if url:
+            result = load_model(None, device, model_url=url)
+        else:
+            result = load_model(model_name, device)
         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_model_button.click(on_load_model_click, inputs=[model_dropdown, device_choice, model_url], outputs=[model_status])
+    generate_button.click(generate_music, inputs=[prompt, seed, cfg_scale, steps, duration, device_choice], outputs=[output_status, output_audio])
 
-    # Load default model on startup
-    iface.load(lambda: on_load_model_click(default_model), inputs=None, outputs=None)
+    # Load default model and resources on startup
+    iface.load(lambda: (load_resources("cpu"), on_load_model_click(default_model, "cpu", None)), inputs=None, outputs=None)
 
 # Launch the interface
 iface.launch()