File size: 1,293 Bytes
df1fa7f | 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 | import io, time
from PIL import Image
from fastapi.testclient import TestClient
from app import app
c = TestClient(app)
colors = ["red","green","blue","yellow","purple"]
def png(col):
b = io.BytesIO(); Image.new("RGB", (640, 480), col).save(b, "PNG"); return b.getvalue()
r = c.post("/generate", data={"prompt": "What color is this image?", "max_new_tokens": 20},
files=[("images", ("a.png", png("red"), "image/png"))])
print("1 img:", r.status_code, r.json())
t = time.time()
files = [("images", (f"{i}.png", png(colors[i % 5]), "image/png")) for i in range(20)]
r = c.post("/generate", data={"prompt": "List the colors of these images.", "max_new_tokens": 60}, files=files)
print("20 img:", r.status_code, r.json(), f"{time.time()-t:.1f}s")
assert r.status_code == 200
files.append(("images", ("x.png", png("red"), "image/png")))
r = c.post("/generate", data={"prompt": "hi"}, files=files)
print("21 img:", r.status_code); assert r.status_code == 400
r = c.post("/generate", data={"prompt": "hi"}, files=[("images", ("bad.png", b"notimage", "image/png"))])
print("bad img:", r.status_code); assert r.status_code == 400
r = c.post("/generate", data={"prompt": "Say hello.", "max_new_tokens": 10})
print("text only:", r.status_code, r.json()); assert r.status_code == 200
|