MCP request flow fixes for URLs

app.py

@@ -28,12 +28,12 @@ import librosa
 import modules.user_history
 from modules.version_info import versions_html, commit_hash, get_xformers_version
 from modules.gradio import *
-from modules.file_utils import get_file_parts, get_filename_from_filepath, convert_title_to_filename, get_unique_file_path, delete_file, download_and_save_image
+from modules.file_utils import get_file_parts, get_filename_from_filepath, convert_title_to_filename, get_unique_file_path, delete_file, download_and_save_image, download_and_save_file
 
 # Added for MCP call
 from smolagents.mcp_client import MCPClient
 from modules.storage import upload_files_to_repo
-from modules.constants import HF_REPO_ID, umg_mcp_server  
+from modules.constants import HF_REPO_ID, umg_mcp_server, TMPDIR, HF_API_TOKEN
 
 MODEL = None
 MODELS = None
@@ -463,6 +463,11 @@ def predict(
         if not isinstance(seed_used, (int, float)):  # Allow float for seed then cast later
             raise gr.Error(f"MCP tool returned a non-numeric seed. Received type: {type(seed_used)}, value: {seed_used}")
 
+        if isinstance(waveform_video_path, str) and (waveform_video_path.startswith("http://") or waveform_video_path.startswith("https://")):
+            waveform_video_path = str(download_and_save_file(waveform_video_path, Path(TMPDIR) / str(profile_username_to_send), HF_API_TOKEN))
+        if isinstance(wave_file_path, str) and (wave_file_path.startswith("http://") or wave_file_path.startswith("https://")):
+            wave_file_path = str(download_and_save_file(wave_file_path, Path(TMPDIR) / str(profile_username_to_send), HF_API_TOKEN))
+
         return waveform_video_path, wave_file_path, int(seed_used)
 
     except Exception as e:

modules/constants.py

@@ -46,7 +46,14 @@ MAX_SEED = np.iinfo(np.int32).max
 TARGET_SIZE = (2688,1536)
 BASE_HEIGHT = 640
 SCALE_FACTOR = (12/5)
-TMPDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+try:
+    if os.environ['TMPDIR']:
+        TMPDIR = os.environ['TMPDIR']
+    else:
+        TMPDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+except:
+    TMPDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+
 os.makedirs(TMPDIR, exist_ok=True)
 
 

modules/file_utils.py

@@ -2,6 +2,10 @@
 import os
 import shutil
 from pathlib import Path
+import requests
+from PIL import Image
+from io import BytesIO
+from urllib.parse import urlparse
 
 def get_file_parts(file_path: str):
     # Split the path into directory and filename
@@ -149,4 +153,37 @@ def download_and_save_image(url: str, dst_folder: Path, token: str = None) -> Pa
     dst = Path(unique_filepath_str)
     dst_folder.mkdir(parents=True, exist_ok=True)
     pil_image.save(dst)
+    return dst
+
+def download_and_save_file(url: str, dst_folder: Path, token: str = None) -> Path:
+    """
+    Downloads a binary file (e.g., audio or video) from a URL with authentication if a token is provided,
+    and saves it in dst_folder with a unique filename.
+
+    Args:
+        url (str): The file URL.
+        dst_folder (Path): The destination folder for the file.
+        token (str, optional): A valid Bearer token.
+
+    Returns:
+        Path: The saved file's path.
+    """
+    headers = {}
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+    
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+    
+    parsed_url = urlparse(url)
+    original_filename = os.path.basename(parsed_url.path)
+    base, ext = os.path.splitext(original_filename)
+    
+    unique_filepath_str = get_unique_file_path(str(dst_folder), base, ext)
+    dst = Path(unique_filepath_str)
+    dst_folder.mkdir(parents=True, exist_ok=True)
+    
+    with open(dst, "wb") as f:
+        f.write(response.content)
+    
     return dst