diff --git a/integration-tests/models/test_idefics.py b/integration-tests/models/test_idefics.py index e49e9c12..80bd5eea 100644 --- a/integration-tests/models/test_idefics.py +++ b/integration-tests/models/test_idefics.py @@ -1,4 +1,5 @@ import pytest +import base64 @pytest.fixture(scope="module") @@ -14,11 +15,17 @@ async def idefics(idefics_handle): await idefics_handle.health(300) return idefics_handle.client +def get_chicken(): + with open("integrations-tests/images/chicken_on_money.png", "rb") as image_file: + encoded_string = base64.b64encode(image_file.read()) + return f"data:image/png;base64,{encoded_string}" + @pytest.mark.asyncio async def test_idefics(idefics, response_snapshot): + chicken = get_chicken() response = await idefics.generate( - "User:![](https://huggingface.co/spaces/HuggingFaceM4/idefics_playground/resolve/6af279110440b7ba0c43f6a0be93b9926f062c42/example_images/chicken_on_money.png)Can you tell me a very short story based on the image?", + f"User:![]({chicken})Can you tell me a very short story based on the image?", max_new_tokens=10, decoder_input_details=True, ) @@ -31,7 +38,7 @@ async def test_idefics(idefics, response_snapshot): async def test_idefics_load(idefics, generate_load, response_snapshot): responses = await generate_load( idefics, - "User:![](https://huggingface.co/spaces/HuggingFaceM4/idefics_playground/resolve/6af279110440b7ba0c43f6a0be93b9926f062c42/example_images/chicken_on_money.png)Can you tell me a very short story based on the image?", + f"User:![]({chicken})Can you tell me a very short story based on the image?", max_new_tokens=10, n=4, )