File size: 4,418 Bytes
52a6059 | 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 | import base64
from io import BytesIO
from unittest.mock import patch
import pytest
from fastapi.testclient import TestClient
from PIL import Image
from flux3_engine import Engine, GenerationRequest, LegacyRequest, Settings
from web import create_app
class FakeEngine:
def __init__(self):
self.pipe = True
self.requests = []
def health(self):
return {'ready': True, 'quantization': 'bf16', 'default_steps': 16}
def generate(self, request, image_format='PNG'):
self.requests.append(request)
data = BytesIO()
Image.new('RGB', (request.width, request.height)).save(data, format=image_format)
return {'image': base64.b64encode(data.getvalue()).decode(), 'metadata': {},
'mime_type': 'image/jpeg' if image_format == 'JPEG' else 'image/png'}
legacy = Engine.legacy
@pytest.fixture
def api(monkeypatch):
monkeypatch.setenv('FLUX_API_TOKEN', 'test-only')
engine = FakeEngine()
with TestClient(create_app(engine)) as client:
yield client, engine
def test_auth_required(api):
client, engine = api
for route in ('/sdxl', '/generate'):
assert client.post(route, json={'inputs': 'portrait'}).status_code == 401
assert client.get('/health').status_code == 401
assert not engine.requests
@pytest.mark.parametrize('route', ['/docs', '/docs/oauth2-redirect', '/redoc', '/openapi.json'])
def test_documentation_disabled(api, route):
client, engine = api
assert client.get(route).status_code == 404
assert client.get(route, headers={'Authorization': 'Bearer test-only'}).status_code == 404
assert not engine.requests
@pytest.mark.parametrize('startup_error', [None, 'Model load failed'])
def test_unready_engine_returns_503(monkeypatch, startup_error):
monkeypatch.setenv('FLUX_API_TOKEN', 'test-only')
engine = FakeEngine()
engine.pipe = True if startup_error else None
with TestClient(create_app(engine, startup_error)) as client:
headers = {'Authorization': 'Bearer test-only'}
assert client.get('/health', headers=headers).status_code == 503
for route in ('/sdxl', '/generate'):
assert client.post(route, headers=headers, json={'inputs': 'portrait'}).status_code == 503
assert not engine.requests
def test_legacy_payload_returns_jpeg_at_16_steps(api):
client, engine = api
response = client.post('/sdxl', headers={'Authorization': 'Bearer test-only'}, json={
'inputs': 'portrait', 'num_inference_steps': 40, 'guidance_scale': 6,
'negative_prompt': 'blurry', 'use_refiner': False, 'high_noise_frac': .8,
'width': 512, 'height': 768})
assert response.status_code == 200
image = Image.open(BytesIO(base64.b64decode(response.json()['image'])))
assert image.format == 'JPEG' and image.size == (512, 768)
assert engine.requests[0].num_inference_steps == 16
def test_native_resolution_and_steps(api):
client, engine = api
headers = {'Authorization': 'Bearer test-only'}
assert client.post('/generate', headers=headers,
json={'inputs': 'portrait', 'width': 768, 'height': 768}).status_code == 200
assert engine.requests[-1].num_inference_steps == 16
for payload in ({'inputs': ' '}, {'inputs': 'p', 'width': 1000},
{'inputs': 'p', 'num_inference_steps': 40}):
assert client.post('/generate', headers=headers, json=payload).status_code == 422
def test_fail_closed_without_api_token(monkeypatch):
monkeypatch.delenv('FLUX_API_TOKEN', raising=False)
with pytest.raises(RuntimeError):
create_app(FakeEngine())
def test_private_repo_does_not_reuse_upstream_commit(monkeypatch):
monkeypatch.setenv('FLUX_MODEL_ID', 'example/private-model')
monkeypatch.delenv('FLUX_MODEL_REVISION', raising=False)
monkeypatch.delenv('FLUX_QUANTIZATION', raising=False)
assert Settings.from_env() == Settings('example/private-model', 'main', 'q8_0')
def test_hf_handler_uses_mounted_repository():
with patch('handler.Engine') as engine:
from handler import EndpointHandler
handler = EndpointHandler('/repository')
payload = {'inputs': 'portrait'}
handler(payload)
engine.assert_called_once_with(model_path='/repository')
engine.return_value.load.assert_called_once()
engine.return_value.legacy.assert_called_once_with(payload)
|