Fixed idefics2 integration tests. (Less VRAM is crucial).

integration-tests/models/__snapshots__/test_idefics2/test_flash_idefics2_next_all_params.json

@@ -0,0 +1,89 @@
+{
+  "details": {
+    "best_of_sequences": null,
+    "finish_reason": "length",
+    "generated_tokens": 10,
+    "prefill": [
+      {
+        "id": 1,
+        "logprob": null,
+        "text": "<s>"
+      },
+      {
+        "id": 3735,
+        "logprob": -8.5625,
+        "text": "Test"
+      },
+      {
+        "id": 2159,
+        "logprob": -10.78125,
+        "text": "request"
+      }
+    ],
+    "seed": 0,
+    "tokens": [
+      {
+        "id": 288,
+        "logprob": -0.2854004,
+        "special": false,
+        "text": "ing"
+      },
+      {
+        "id": 264,
+        "logprob": -0.37573242,
+        "special": false,
+        "text": " a"
+      },
+      {
+        "id": 633,
+        "logprob": -0.09301758,
+        "special": false,
+        "text": " new"
+      },
+      {
+        "id": 4480,
+        "logprob": -0.3322754,
+        "special": false,
+        "text": " feature"
+      },
+      {
+        "id": 297,
+        "logprob": -0.8510742,
+        "special": false,
+        "text": " in"
+      },
+      {
+        "id": 272,
+        "logprob": -0.13464355,
+        "special": false,
+        "text": " the"
+      },
+      {
+        "id": 2039,
+        "logprob": 0.0,
+        "special": false,
+        "text": " game"
+      },
+      {
+        "id": 28723,
+        "logprob": -0.89990234,
+        "special": false,
+        "text": "."
+      },
+      {
+        "id": 13,
+        "logprob": 0.0,
+        "special": false,
+        "text": "\n"
+      },
+      {
+        "id": 13,
+        "logprob": 0.0,
+        "special": false,
+        "text": "\n"
+      }
+    ],
+    "top_tokens": null
+  },
+  "generated_text": "Test requesting a new feature in the game.\n\n"
+}

integration-tests/models/__snapshots__/test_idefics2/test_flash_idefics2_next_simple.json

@@ -0,0 +1,73 @@
+{
+  "details": {
+    "best_of_sequences": null,
+    "finish_reason": "length",
+    "generated_tokens": 10,
+    "prefill": [],
+    "seed": null,
+    "tokens": [
+      {
+        "id": 330,
+        "logprob": -0.13000488,
+        "special": false,
+        "text": " A"
+      },
+      {
+        "id": 13088,
+        "logprob": -0.6713867,
+        "special": false,
+        "text": " chicken"
+      },
+      {
+        "id": 349,
+        "logprob": -0.2980957,
+        "special": false,
+        "text": " is"
+      },
+      {
+        "id": 6398,
+        "logprob": -0.060638428,
+        "special": false,
+        "text": " sitting"
+      },
+      {
+        "id": 356,
+        "logprob": -0.27319336,
+        "special": false,
+        "text": " on"
+      },
+      {
+        "id": 264,
+        "logprob": -0.140625,
+        "special": false,
+        "text": " a"
+      },
+      {
+        "id": 17972,
+        "logprob": -0.040405273,
+        "special": false,
+        "text": " pile"
+      },
+      {
+        "id": 302,
+        "logprob": -0.0002708435,
+        "special": false,
+        "text": " of"
+      },
+      {
+        "id": 2445,
+        "logprob": -0.095336914,
+        "special": false,
+        "text": " money"
+      },
+      {
+        "id": 28723,
+        "logprob": -0.0068359375,
+        "special": false,
+        "text": "."
+      }
+    ],
+    "top_tokens": null
+  },
+  "generated_text": " A chicken is sitting on a pile of money."
+}

integration-tests/models/test_idefics2.py

@@ -0,0 +1,81 @@
+import pytest
+import base64
+
+
+# TODO fix the server parsser to count inline image tokens correctly
+def get_chicken():
+    with open("integration-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.decode('utf-8')}"
+
+
+@pytest.fixture(scope="module")
+def flash_idefics2_next_handle(launcher):
+    with launcher(
+        "HuggingFaceM4/idefics2-8b",
+    ) as handle:
+        yield handle
+
+
+@pytest.fixture(scope="module")
+async def flash_idefics2_next(flash_idefics2_next_handle):
+    await flash_idefics2_next_handle.health(300)
+    return flash_idefics2_next_handle.client
+
+
+@pytest.mark.asyncio
+@pytest.mark.private
+async def test_flash_idefics2_next_simple(flash_idefics2_next, response_snapshot):
+    chicken = get_chicken()
+    response = await flash_idefics2_next.generate(
+        f"User:![]({chicken})Write me a short story<end_of_utterance> \nAssistant:",
+        max_new_tokens=10,
+    )
+    assert (
+        response.generated_text == " A chicken is sitting on a pile of money."
+    ), f"{repr(response.generated_text)}"
+    assert response.details.generated_tokens == 10
+    assert response == response_snapshot
+
+
+@pytest.mark.asyncio
+@pytest.mark.private
+async def test_flash_idefics2_next_all_params(flash_idefics2_next, response_snapshot):
+    response = await flash_idefics2_next.generate(
+        "Test request",
+        max_new_tokens=10,
+        repetition_penalty=1.2,
+        return_full_text=True,
+        stop_sequences=["test"],
+        temperature=0.5,
+        top_p=0.9,
+        top_k=10,
+        truncate=5,
+        typical_p=0.9,
+        watermark=True,
+        decoder_input_details=True,
+        seed=0,
+    )
+
+    assert response.details.generated_tokens == 10
+    assert response == response_snapshot
+
+
+@pytest.mark.asyncio
+@pytest.mark.private
+async def test_flash_idefics2_next_load(
+    flash_idefics2_next, generate_load, response_snapshot
+):
+    chicken = get_chicken()
+    responses = await generate_load(
+        flash_idefics2_next,
+        f"User:![]({chicken})Write me a short story<end_of_utterance> \nAssistant:",
+        max_new_tokens=10,
+        n=4,
+    )
+    generated_texts = [r.generated_text for r in responses]
+    assert generated_texts[0] == " A chicken is sitting on a pile of money."
+    assert len(generated_texts) == 4
+    assert all([r.generated_text == generated_texts[0] for r in responses])
+
+    assert responses == response_snapshot