154 lines
4.6 KiB
Python
154 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
测试下载进度API的脚本
|
|
"""
|
|
|
|
import requests
|
|
import json
|
|
import time
|
|
|
|
# ComfyUI服务器地址
|
|
BASE_URL = "http://localhost:8188"
|
|
|
|
def test_download_model():
|
|
"""测试开始下载模型"""
|
|
url = f"{BASE_URL}/api/download-model"
|
|
|
|
data = {
|
|
"model_id": "damo/sd-v1-5",
|
|
"model_type": "checkpoints",
|
|
"dest_dir": None # 使用默认目录
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, json=data)
|
|
result = response.json()
|
|
|
|
if result.get("success"):
|
|
download_id = result["data"]["download_id"]
|
|
print(f"✅ 下载已开始，下载ID: {download_id}")
|
|
return download_id
|
|
else:
|
|
print(f"❌ 下载启动失败: {result.get('message')}")
|
|
return None
|
|
|
|
except Exception as e:
|
|
print(f"❌ 请求失败: {e}")
|
|
return None
|
|
|
|
def poll_download_progress(download_id):
|
|
"""轮询下载进度"""
|
|
url = f"{BASE_URL}/api/download-progress/{download_id}"
|
|
|
|
print(f"\n📊 开始监控下载进度 (ID: {download_id})")
|
|
print("=" * 60)
|
|
|
|
while True:
|
|
try:
|
|
response = requests.get(url)
|
|
result = response.json()
|
|
|
|
if result.get("success"):
|
|
progress_data = result["data"]
|
|
status = progress_data["status"]
|
|
percentage = progress_data["percentage"]
|
|
progress = progress_data["progress"]
|
|
file_size = progress_data["file_size"]
|
|
speed = progress_data["speed"]
|
|
estimated_time = progress_data["estimated_time"]
|
|
|
|
# 格式化显示
|
|
if file_size > 0:
|
|
progress_mb = progress / (1024 * 1024)
|
|
total_mb = file_size / (1024 * 1024)
|
|
progress_str = f"{progress_mb:.1f}MB / {total_mb:.1f}MB"
|
|
else:
|
|
progress_str = f"{progress} bytes"
|
|
|
|
speed_str = f"{speed:.2f} B/s" if speed else "N/A"
|
|
eta_str = f"{estimated_time:.1f}s" if estimated_time else "N/A"
|
|
|
|
print(f"\r📥 {status.upper()}: {percentage:.1f}% | {progress_str} | 速度: {speed_str} | 剩余时间: {eta_str}", end="", flush=True)
|
|
|
|
if status in ["completed", "failed"]:
|
|
print() # 换行
|
|
if status == "completed":
|
|
print(f"✅ 下载完成！总用时: {progress_data.get('total_time', 'N/A')}s")
|
|
else:
|
|
print(f"❌ 下载失败: {progress_data.get('error_message', 'Unknown error')}")
|
|
break
|
|
|
|
else:
|
|
print(f"\n❌ 获取进度失败: {result.get('message')}")
|
|
break
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ 请求失败: {e}")
|
|
break
|
|
|
|
time.sleep(1) # 每秒轮询一次
|
|
|
|
def list_all_downloads():
|
|
"""列出所有下载"""
|
|
url = f"{BASE_URL}/api/download-progress"
|
|
|
|
try:
|
|
response = requests.get(url)
|
|
result = response.json()
|
|
|
|
if result.get("success"):
|
|
downloads = result["data"]["downloads"]
|
|
count = result["data"]["count"]
|
|
print(f"\n📋 当前活跃下载: {count} 个")
|
|
|
|
for download_id in downloads:
|
|
print(f" - {download_id}")
|
|
else:
|
|
print(f"❌ 获取下载列表失败: {result.get('message')}")
|
|
|
|
except Exception as e:
|
|
print(f"❌ 请求失败: {e}")
|
|
|
|
def clear_download_progress(download_id):
|
|
"""清理下载进度记录"""
|
|
url = f"{BASE_URL}/api/download-progress/{download_id}"
|
|
|
|
try:
|
|
response = requests.delete(url)
|
|
result = response.json()
|
|
|
|
if result.get("success"):
|
|
print(f"✅ 已清理下载进度记录: {download_id}")
|
|
else:
|
|
print(f"❌ 清理失败: {result.get('message')}")
|
|
|
|
except Exception as e:
|
|
print(f"❌ 请求失败: {e}")
|
|
|
|
def main():
|
|
"""主函数"""
|
|
print("🚀 下载进度API测试")
|
|
print("=" * 60)
|
|
|
|
# 1. 开始下载
|
|
download_id = test_download_model()
|
|
if not download_id:
|
|
return
|
|
|
|
# 2. 列出所有下载
|
|
list_all_downloads()
|
|
|
|
# 3. 监控下载进度
|
|
poll_download_progress(download_id)
|
|
|
|
# 4. 清理进度记录（可选）
|
|
# clear_download_progress(download_id)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
|
|
|
|
|
|
|