Spaces:
Sleeping
Sleeping
| """The ZeroGPU Space: transformers VLM backend and the Gradio app, driven over its API with gradio_client. | |
| Uses a tiny random Qwen3-VL built on disk (tests/tiny_vlm.py) so nothing is downloaded; set | |
| MRAG_VLM_ID to a real model (e.g. Qwen/Qwen3-VL-2B-Instruct) to run the same tests against it. | |
| """ | |
| import os | |
| import socket | |
| import sys | |
| from pathlib import Path | |
| import pytest | |
| from conftest import ROOT | |
| pytest.importorskip("transformers") | |
| def vlm_id(tmp_path_factory): | |
| if os.environ.get("MRAG_VLM_ID"): | |
| return os.environ["MRAG_VLM_ID"] | |
| try: | |
| from transformers import Qwen3VLForConditionalGeneration # noqa: F401 | |
| except ImportError: | |
| pytest.skip("transformers without Qwen3-VL (needs >=4.57)") | |
| sys.path.insert(0, str(ROOT / "tests")) | |
| import tiny_vlm | |
| return str(tiny_vlm.build(tmp_path_factory.mktemp("tiny_qwen3vl"))) | |
| def vlm(vlm_id): | |
| from mrag.hf_vlm import HFVLM | |
| return HFVLM(vlm_id) | |
| def test_vlm_image_tokens_match_vision_grid(vlm): | |
| from PIL import Image | |
| ims = [Image.new("RGB", (900, 500), "red"), Image.new("RGB", (300, 600), "blue")] | |
| inp = vlm.build_inputs("system", "Question: what colour?", ims) | |
| grid = inp["image_grid_thw"] | |
| merge = vlm.processor.image_processor.merge_size | |
| n_img_tokens = int((inp["input_ids"] == vlm.model.config.image_token_id).sum()) | |
| assert grid.shape[0] == 2 and n_img_tokens == int(grid.prod(-1).sum()) // merge ** 2 | |
| def test_vlm_streams_text(vlm): | |
| from PIL import Image | |
| pieces = list(vlm.stream("system", "Question: hi", [Image.new("RGB", (64, 64))], max_new_tokens=8)) | |
| assert pieces and all(isinstance(p, str) for p in pieces) | |
| def test_vlm_errors_are_raised_not_hung(vlm): | |
| with pytest.raises(Exception): | |
| list(vlm.stream("system", "x", [], max_new_tokens=-5)) | |
| def test_generator_transformers_backend(vlm, tmp_path): | |
| from PIL import Image | |
| from mrag.config import Config | |
| from mrag.generate import Generator, ImageHit, TextHit | |
| p = tmp_path / "a.png" | |
| Image.new("RGB", (3000, 1500), "green").save(p) | |
| g = Generator(Config(llm_backend="transformers", max_new_tokens=6, max_image_side=256), vlm=vlm) | |
| out = g.answer("q?", [TextHit("1", "fact", "/d/x.pdf", 1, "text", 1.0)], | |
| [ImageHit("i", str(p), "/d/x.pdf", 1, "", 0.3)]) | |
| assert isinstance(out, str) | |
| def _free_port(): | |
| s = socket.socket() | |
| s.bind(("127.0.0.1", 0)) | |
| port = s.getsockname()[1] | |
| s.close() | |
| return port | |
| def client(vlm_id, tmp_path_factory): | |
| pytest.importorskip("gradio_client") | |
| os.environ["MRAG_VLM_ID"] = vlm_id | |
| os.environ.setdefault("MRAG_TEXT_BACKEND", "wordllama") | |
| os.chdir(tmp_path_factory.mktemp("space")) # demo_docs is built next to app.py; keep tests clean | |
| sys.path.insert(0, str(ROOT)) | |
| import app | |
| from gradio_client import Client | |
| port = _free_port() | |
| app.demo.queue().launch(server_port=port, prevent_thread_lock=True, quiet=True) | |
| yield Client(f"http://127.0.0.1:{port}/", verbose=False), app | |
| app.demo.close() | |
| def test_app_demo_docs_and_streamed_answer(client): | |
| c, app = client | |
| status = c.predict(api_name="/load_demo") | |
| assert "7 file(s)" in status and "report.pdf" in status | |
| chat, gallery, sources, box, *_ = c.predict( | |
| "How long is the warranty on the Nimbus R2?", [], 6, 3, True, | |
| api_name="/prepare_answer") | |
| assert len(chat) == 2 and chat[0]["role"] == "user" | |
| assert "36-month" in sources and "report.pdf p.3" in sources | |
| assert gallery and "report.pdf p.3" in gallery[0]["caption"] | |
| job = c.submit(64, chat, api_name="/vlm_answer") | |
| outputs = list(job) | |
| final = job.result() | |
| assert outputs and final[-1]["role"] == "assistant" and final[-1]["content"] | |
| assert box["value"] == "" and box["interactive"] is False | |
| # gradio_client calls endpoints individually; the browser runs this chained event. | |
| c.predict(api_name="/finish_answer") | |
| def test_app_upload_and_session_isolation(client, corpus): | |
| from gradio_client import Client, handle_file | |
| c, _ = client | |
| fresh = Client(c.src, verbose=False) | |
| with pytest.raises(Exception, match="Index some files"): | |
| fresh.predict("a cat?", [], 6, 3, True, api_name="/prepare_answer") | |
| status, _ = fresh.predict([handle_file(str(corpus / "images" / "office_cat.jpg"))], api_name="/index_files") | |
| assert "1 file(s)" in status and "office_cat.jpg" in status | |
| chat, gallery, sources, *_ = fresh.predict("show me a picture of a cat", [], 6, 3, True, api_name="/prepare_answer") | |
| assert gallery and "office_cat.jpg" in gallery[0]["caption"] | |
| assert "report.pdf" not in sources | |
| final = fresh.predict(64, chat, api_name="/vlm_answer") | |
| assert final[-1]["role"] == "assistant" | |
| fresh.predict(api_name="/finish_answer") | |
| def test_app_busy_guard_and_failure_cleanup(client): | |
| c, app = client | |
| chat, *_ = c.predict("Who is the CEO?", [], 6, 0, False, api_name="/prepare_answer") | |
| with pytest.raises(Exception, match="still running"): | |
| c.predict(api_name="/reset") | |
| # This handler also runs when ZeroGPU rejects a request or generation fails. | |
| c.predict(api_name="/finish_answer") | |
| with pytest.raises(Exception, match="Submit a question"): | |
| c.predict(64, chat, api_name="/vlm_answer") | |
| c.predict(api_name="/finish_answer") | |
| def test_app_cumulative_upload_limit(client, tmp_path): | |
| from gradio_client import Client, handle_file | |
| c, app = client | |
| fresh = Client(c.src, verbose=False) | |
| paths = [] | |
| for i in range(app.MAX_FILES + 1): | |
| p = tmp_path / f"file{i}.txt" | |
| p.write_text(f"Document {i}: the robot payload is 25 kilograms.") | |
| paths.append(handle_file(str(p))) | |
| status, _ = fresh.predict(paths[:app.MAX_FILES], api_name="/index_files") | |
| assert f"{app.MAX_FILES} file(s)" in status | |
| with pytest.raises(Exception, match="indexed files per session"): | |
| fresh.predict(paths[-1:], api_name="/index_files") | |
| fresh.predict(api_name="/reset") | |
| def test_app_generation_is_success_only_and_serialized(client): | |
| _, app = client | |
| gpu_id = next(i for i, fn in app.demo.fns.items() if fn.fn == app.vlm_answer) | |
| binding = next(d for d in app.demo.config["dependencies"] if d["id"] == gpu_id) | |
| assert binding["trigger_only_on_success"] is True | |
| assert app.demo.fns[gpu_id].concurrency_limit == 1 | |