File size: 4,916 Bytes
7906fb3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import io
import json
import uuid
import urllib.request
import urllib.parse
import websocket
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
# --- Cấu hình ---
# Địa chỉ máy chủ ComfyUI, mặc định là localhost
COMFYUI_SERVER_ADDRESS = "127.0.0.1:8188"
# ID của client, dùng để nhận diện trong WebSocket
CLIENT_ID = str(uuid.uuid4())
# Khởi tạo ứng dụng FastAPI
app = FastAPI(title="Clarity AI Upscaler API")
def get_image(filename, subfolder, folder_type):
"""Lấy file ảnh từ máy chủ ComfyUI."""
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
url_values = urllib.parse.urlencode(data)
with urllib.request.urlopen(f"http://{COMFYUI_SERVER_ADDRESS}/view?{url_values}") as response:
return response.read()
def queue_prompt(prompt_workflow):
"""Gửi yêu cầu thực thi quy trình đến ComfyUI."""
p = {"prompt": prompt_workflow, "client_id": CLIENT_ID}
data = json.dumps(p).encode('utf-8')
req = urllib.request.Request(f"http://{COMFYUI_SERVER_ADDRESS}/prompt", data=data)
return json.loads(urllib.request.urlopen(req).read())
def upload_image(image_bytes: bytes, filename: str = "input_image.png"):
"""Tải ảnh lên thư mục input của ComfyUI."""
import requests
files = {"image": (filename, image_bytes, 'image/png'), "overwrite": (None, 'true')}
response = requests.post(f"http://{COMFYUI_SERVER_ADDRESS}/upload/image", files=files)
if response.status_code == 200:
return response.json()
else:
raise HTTPException(status_code=500, detail=f"Failed to upload image: {response.text}")
def get_history(prompt_id):
"""Lấy lịch sử thực thi của một prompt."""
with urllib.request.urlopen(f"http://{COMFYUI_SERVER_ADDRESS}/history/{prompt_id}") as response:
return json.loads(response.read())
def track_execution_and_get_output(prompt_id):
"""Theo dõi tiến trình qua WebSocket và lấy ảnh kết quả."""
ws_url = f"ws://{COMFYUI_SERVER_ADDRESS}/ws?clientId={CLIENT_ID}"
ws = websocket.WebSocket()
ws.connect(ws_url)
while True:
out = ws.recv()
if isinstance(out, str):
message = json.loads(out)
if message['type'] == 'executing':
data = message['data']
if data['node'] is None and data['prompt_id'] == prompt_id:
break # Đã thực thi xong
else:
continue # Bỏ qua các tin nhắn không phải dạng text
ws.close()
history = get_history(prompt_id)[prompt_id]
for node_id, node_output in history['outputs'].items():
if 'images' in node_output:
for image in node_output['images']:
if image['type'] == 'output':
image_data = get_image(image['filename'], image['subfolder'], image['type'])
return image_data
raise HTTPException(status_code=500, detail="Không tìm thấy ảnh kết quả.")
@app.post("/upscale/",
summary="Nâng cấp và tinh chỉnh hình ảnh",
response_description="Hình ảnh đã được xử lý ở định dạng PNG")
async def upscale_image(file: UploadFile = File(..., description="File ảnh cần xử lý.")):
"""
Nhận một file ảnh, xử lý nó thông qua quy trình ComfyUI, và trả về kết quả.
"""
try:
# 1. Đọc file workflow từ định dạng API
with open("workflow_api.json", "r", encoding="utf-8") as f:
prompt_workflow = json.load(f)
# 2. Đọc và tải ảnh đầu vào lên ComfyUI
image_bytes = await file.read()
upload_response = upload_image(image_bytes)
input_filename = upload_response['name']
# 3. Tìm node LoadImage (ID=17) và cập nhật tên file
# LƯU Ý: ID '17' là ID của node LoadImage trong workflow của bạn.
# Nếu bạn thay đổi workflow, ID này có thể thay đổi.
prompt_workflow["17"]["inputs"]["image"] = input_filename
# 4. Gửi quy trình để thực thi
queue_response = queue_prompt(prompt_workflow)
prompt_id = queue_response['prompt_id']
# 5. Theo dõi tiến trình và lấy ảnh kết quả
result_image_bytes = track_execution_and_get_output(prompt_id)
# 6. Trả về ảnh kết quả
return StreamingResponse(io.BytesIO(result_image_bytes), media_type="image/png")
except FileNotFoundError:
raise HTTPException(status_code=500, detail="Không tìm thấy file workflow_api.json.")
except Exception as e:
raise HTTPException(status_code=500, detail=f"Đã xảy ra lỗi: {str(e)}")
# Để chạy máy chủ này, mở terminal và gõ lệnh:
# uvicorn api_server:app --reload
|