inpaintServer/tests/load_test_lama.py

76 lines
2.7 KiB
Python

import asyncio
import httpx
import base64
from pathlib import Path
import time
# 테스트 설정
BASE_URL = "http://localhost:8008/api/v1/inpaint"
# 테스트에 사용할 이미지와 마스크 경로
IMAGE_PATH = Path(__file__).parent / "test_image.png"
MASK_PATH = Path(__file__).parent / "test_mask.png"
CONCURRENT_REQUESTS = 8 # 동시에 보낼 요청 수 (min_sessions < 요청 수 < max_sessions * 2 정도로 설정)
def encode_image(image_path: Path) -> str:
"""이미지 파일을 base64로 인코딩"""
try:
with open(image_path, "rb") as f:
return base64.b64encode(f.read()).decode('utf-8')
except FileNotFoundError as e:
print(f"오류: 테스트 이미지를 찾을 수 없습니다 - {e}")
print("먼저 tests/generate_test_images.py 를 실행하여 테스트 이미지를 생성해주세요.")
raise
async def send_request(client: httpx.AsyncClient, request_num: int, image_b64: str, mask_b64: str):
"""단일 API 요청을 보냄"""
payload = {
"image": image_b64,
"mask": mask_b64,
"model_name": "simple-lama",
"ldm_steps": 20,
"ldm_sampler": "plms",
"hd_strategy": "Original",
"hd_strategy_crop_margin": 128,
"hd_strategy_crop_trigger_size": 800,
"hd_strategy_resize_limit": 1280
}
print(f"[{request_num:02d}] 요청 시작...")
start_time = time.time()
try:
response = await client.post(BASE_URL, json=payload, timeout=120)
duration = time.time() - start_time
if response.status_code == 200:
print(f"✅ [{request_num:02d}] 요청 성공 ({duration:.2f}초)")
# with open(f"result_{request_num}.png", "wb") as f:
# f.write(response.content)
else:
print(f"❌ [{request_num:02d}] 요청 실패 - 상태 코드: {response.status_code} ({duration:.2f}초)")
print(f" 응답: {response.text[:200]}")
except httpx.ReadTimeout:
duration = time.time() - start_time
print(f"❌ [{request_num:02d}] 요청 시간 초과 ({duration:.2f}초)")
except httpx.ConnectError as e:
print(f"❌ [{request_num:02d}] 연결 실패: {e}")
async def main():
"""메인 실행 함수"""
print("SimpleLama 부하 테스트 시작...")
print(f"동시 요청 수: {CONCURRENT_REQUESTS}")
# 이미지 파일 인코딩
image_b64 = encode_image(IMAGE_PATH)
mask_b64 = encode_image(MASK_PATH)
async with httpx.AsyncClient() as client:
tasks = [
send_request(client, i, image_b64, mask_b64)
for i in range(CONCURRENT_REQUESTS)
]
await asyncio.gather(*tasks)
if __name__ == "__main__":
asyncio.run(main())