Add generation params to DB, fix GUI star layout/refresh, add sample generator script

2026-06-28 16:56:44 +03:00
parent 68cba874d5
commit 09f6f377c7
4 changed files with 223 additions and 60 deletions
@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+"""Generate a sample dataset through the MCP server."""
+
+import asyncio
+import os
+import sys
+
+from mcp.client.session import ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+SAMPLE_SPECS = [
+    {
+        "prompt": "a brave knight in shining armor holding a sword",
+        "output_path": "tetra/knight_01.png",
+        "seed": 42,
+    },
+    {
+        "prompt": "an elven archer with a longbow and green cloak",
+        "output_path": "tetra/archer_01.png",
+        "seed": 43,
+    },
+    {
+        "prompt": "a robed mage casting a blue fireball",
+        "output_path": "tetra/mage_01.png",
+        "seed": 44,
+    },
+    {
+        "prompt": "a peasant worker carrying a wooden hammer",
+        "output_path": "tetra/worker_01.png",
+        "seed": 45,
+    },
+    {
+        "prompt": "a crystal golem with geometric facets",
+        "output_path": "tetra/golem_01.png",
+        "seed": 46,
+    },
+    {
+        "prompt": "a chaos demon with horns and jagged armor",
+        "output_path": "tetra/demon_01.png",
+        "seed": 47,
+    },
+    {
+        "prompt": "an ancient treant with glowing green eyes",
+        "output_path": "tetra/treant_01.png",
+        "seed": 48,
+    },
+    {
+        "prompt": "a skeleton warrior with a rusted sword",
+        "output_path": "tetra/skeleton_01.png",
+        "seed": 49,
+    },
+]
+
+
+async def main():
+    server_path = os.path.join(BASE_DIR, "server.py")
+    params = StdioServerParameters(
+        command=sys.executable,
+        args=[server_path],
+        env={**os.environ, "PYTHONUNBUFFERED": "1"},
+    )
+
+    async with stdio_client(params) as (read, write):
+        async with ClientSession(read, write) as session:
+            await session.initialize()
+
+            print("Calling batch_generate...")
+            result = await session.call_tool("batch_generate", {"specs": SAMPLE_SPECS})
+
+            print("\nResults:")
+            for item in result.content:
+                if item.type == "text":
+                    print(item.text)
+
+            stats = await session.call_tool("db_stats", {})
+            print("\nDB stats:")
+            for item in stats.content:
+                if item.type == "text":
+                    print(item.text)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())