ComfyUI-CustomNode/nodes/file_upload.py

74 lines
2.8 KiB
Python

# -*- coding: utf-8 -*-
"""
Author charon
Date 2025/9/7 13:12
"""
import os
import requests
import mimetypes
import folder_paths
class FileUploadNode:
ENVIRONMENT_URL_MAP = {
"prod": "https://bowongai-prod--text-video-agent-fastapi-app.modal.run",
"test": "https://bowongai-test--text-video-agent-fastapi-app.modal.run",
"dev": "https://bowongai-dev--text-video-agent-fastapi-app.modal.run"
}
API_ENDPOINT = "/api/file/upload/s3"
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"file_to_upload": ("*", {"file_upload": True}),
"environment": (list(s.ENVIRONMENT_URL_MAP.keys()), {"default": "dev"}),
"form_field_name": ("STRING", {"multiline": False, "default": "file"}),
},
}
CATEGORY = "不忘科技-自定义节点🚩/utils/文件上传"
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("url",)
FUNCTION = "handler_file_upload"
def handler_file_upload(self, file_to_upload, environment, form_field_name):
if not file_to_upload or not isinstance(file_to_upload, str):
return (f"Error: Invalid file input. Please upload a file.",)
base_url = self.ENVIRONMENT_URL_MAP.get(environment)
if not base_url:
error_message = f"Error: Invalid environment '{environment}' selected."
return (error_message,)
full_api_url = base_url + self.API_ENDPOINT
local_filepath = folder_paths.get_annotated_filepath(file_to_upload)
if not os.path.exists(local_filepath):
return (f"Error: File not found at path: {local_filepath}",)
headers = {'accept': 'application/json'}
filename = os.path.basename(local_filepath)
mime_type, _ = mimetypes.guess_type(local_filepath)
if mime_type is None:
mime_type = 'application/octet-stream'
try:
with open(local_filepath, 'rb') as f:
files = {form_field_name: (filename, f, mime_type)}
print(f"Uploading '{filename}' to '{full_api_url}'...")
response = requests.post(full_api_url, headers=headers, files=files)
response.raise_for_status()
json_response = response.json()
if json_response.get("status"):
result_url = json_response.get("data")
print(f"Success! URL: {result_url}")
return (result_url,)
else:
api_msg = json_response.get("msg", "Unknown API error")
error_message = f"Error from API: '{api_msg}'"
return (error_message,)
except Exception as e:
raise ValueError(f'File upload process failed: {e}')