kgmodule-utils 0.4.2__tar.gz → 0.4.3__tar.gz

{kgmodule_utils-0.4.2 → kgmodule_utils-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kgmodule-utils
-Version: 0.4.2
+Version: 0.4.3
 Summary: Shared types, graph store, semantic index, and pipeline base for the KGModule SDK
 License: Elastic-2.0
 License-File: LICENSE

{kgmodule_utils-0.4.2 → kgmodule_utils-0.4.3}/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "kgmodule-utils"
-version = "0.4.2"
+version = "0.4.3"
 description = "Shared types, graph store, semantic index, and pipeline base for the KGModule SDK"
 readme = "README.md"
 license = { text = "Elastic-2.0" }

{kgmodule_utils-0.4.2 → kgmodule_utils-0.4.3}/src/kg_utils/__init__.py

@@ -24,4 +24,4 @@ Optional extras
     pip install 'kgmodule-utils[synthesis-mflux]'  # + mflux (Apple Silicon local gen)
 """
 
-__version__ = "0.4.2"
+__version__ = "0.4.3"

{kgmodule_utils-0.4.2 → kgmodule_utils-0.4.3}/src/kg_utils/synthesis/_image.py

@@ -51,6 +51,29 @@ _GPT_IMAGE_SIZES: dict[str, str] = {
 }
 
 
+# ---------------------------------------------------------------------------
+# Size parsing
+# ---------------------------------------------------------------------------
+
+
+def _parse_size(size: str | None) -> tuple[int, int] | None:
+    """Parse an explicit ``"WIDTHxHEIGHT"`` string into an ``(width, height)`` pair.
+
+    :param size: Size string such as ``"768x512"`` (case-insensitive ``x``), or None.
+    :returns: ``(width, height)`` when *size* parses to two positive ints, else None.
+    """
+    if not size:
+        return None
+    try:
+        w_str, h_str = size.lower().split("x", 1)
+        width, height = int(w_str), int(h_str)
+    except (ValueError, AttributeError):
+        return None
+    if width <= 0 or height <= 0:
+        return None
+    return width, height
+
+
 # ---------------------------------------------------------------------------
 # Synthesizer
 # ---------------------------------------------------------------------------
@@ -97,6 +120,7 @@ class ImageSynthesizer:
         seed: int | None = None,
         steps: int | None = None,
         model: str | None = None,
+        size: str | None = None,
     ) -> PILImage:
         """Generate an image and return a PIL Image.
 
@@ -105,6 +129,9 @@ class ImageSynthesizer:
         :param seed: Random seed for reproducibility (random int if omitted).
         :param steps: Override inference steps (mflux backends only; ignored for OpenAI).
         :param model: Override the configured model for this single call.
+        :param size: Explicit ``"WIDTHxHEIGHT"`` override (mflux backends only). When given
+                     and valid, it takes priority over the *aspect_ratio* size table. OpenAI
+                     backends ignore it because they accept only a fixed set of sizes.
         :returns: PIL Image.
         """
         cfg = self._cfg
@@ -118,6 +145,7 @@ class ImageSynthesizer:
                 seed=seed,
                 steps=effective_steps,
                 aspect_ratio=aspect_ratio,
+                size=size,
             )
         elif cfg.backend == ImageBackend.MFLUX_SERVE:
             return self._generate_via_server(
@@ -126,6 +154,7 @@ class ImageSynthesizer:
                 seed=seed,
                 steps=effective_steps,
                 aspect_ratio=aspect_ratio,
+                size=size,
             )
         else:
             return self._generate_openai(
@@ -142,6 +171,7 @@ class ImageSynthesizer:
         seed: int | None = None,
         steps: int | None = None,
         model: str | None = None,
+        size: str | None = None,
     ) -> str:
         """Generate an image and return it as a base64-encoded PNG string.
 
@@ -152,9 +182,12 @@ class ImageSynthesizer:
         :param seed: Random seed for reproducibility.
         :param steps: Override inference steps (mflux backends only).
         :param model: Override the configured model for this single call.
+        :param size: Explicit ``"WIDTHxHEIGHT"`` override (mflux backends only).
         :returns: Base64-encoded PNG string.
         """
-        img = self.generate(prompt, aspect_ratio=aspect_ratio, seed=seed, steps=steps, model=model)
+        img = self.generate(
+            prompt, aspect_ratio=aspect_ratio, seed=seed, steps=steps, model=model, size=size
+        )
         buf = BytesIO()
         img.save(buf, format="PNG")
         return base64.b64encode(buf.getvalue()).decode()
@@ -171,8 +204,9 @@ class ImageSynthesizer:
         seed: int | None,
         steps: int,
         aspect_ratio: str,
+        size: str | None = None,
     ) -> PILImage:
-        width, height = _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
+        width, height = _parse_size(size) or _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
         effective_seed = seed if seed is not None else random.randint(0, 2**31 - 1)
         flux = self._load_mflux(model)
         result = flux.generate_image(
@@ -194,11 +228,12 @@ class ImageSynthesizer:
         seed: int | None,
         steps: int,
         aspect_ratio: str,
+        size: str | None = None,
     ) -> PILImage:
         import httpx
         from PIL import Image  # type: ignore[import-unresolved]
 
-        width, height = _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
+        width, height = _parse_size(size) or _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
         payload: dict[str, Any] = {
             "prompt": prompt,
             "n": 1,

{kgmodule_utils-0.4.2 → kgmodule_utils-0.4.3}/src/kg_utils/worker/client.py

@@ -8,6 +8,7 @@ This module centralizes payload construction and response/error decoding for
 from __future__ import annotations
 
 import json
+from typing import Any, cast
 
 import httpx
 
@@ -29,8 +30,9 @@ def _format_error_data(error_data: object) -> str:
         return str(decoded)
 
     if isinstance(error_data, dict):
-        err_type = error_data.get("error_type", "Unknown")
-        err_msg = error_data.get("error_message", str(error_data))
+        d = cast("dict[str, Any]", error_data)
+        err_type = d.get("error_type", "Unknown")
+        err_msg = d.get("error_message", str(error_data))
         return f"{err_type}: {err_msg}"
 
     return str(error_data)
@@ -41,10 +43,11 @@ def extract_worker_error(data: object) -> str | None:
     if not isinstance(data, dict):
         return str(data)
 
-    if data.get("status") == "FAILED" or "error_type" in data:
-        return _format_error_data(data.get("error", data))
+    d = cast("dict[str, Any]", data)
+    if d.get("status") == "FAILED" or "error_type" in d:
+        return _format_error_data(d.get("error", d))
 
-    out = data.get("output")
+    out = d.get("output")
     if isinstance(out, dict) and isinstance(out.get("error"), str):
         return out["error"]
 
@@ -60,7 +63,8 @@ def decode_worker_response(data: object) -> dict:
     if not isinstance(data, dict):
         raise WorkerError(f"unexpected worker response type: {type(data).__name__}")
 
-    out = data.get("output", data)
+    d2 = cast("dict[str, Any]", data)
+    out = d2.get("output", d2)
     if not isinstance(out, dict):
         raise WorkerError(f"unexpected worker output type: {type(out).__name__}")
     return out
@@ -133,14 +137,18 @@ class WorkerClient:
         image_backend: str = "",
         aspect_ratio: str = "3:2",
         steps: int | None = None,
+        size: str | None = None,
     ) -> tuple[str | None, str | None, str | None, str | None]:
-        payload: dict = {"input": {"op": "imagine", "prompt": prompt, "aspect_ratio": aspect_ratio}}
+        inp: dict[str, Any] = {"op": "imagine", "prompt": prompt, "aspect_ratio": aspect_ratio}
         if image_backend:
-            payload["input"]["image_backend"] = image_backend
+            inp["image_backend"] = image_backend
         if steps is not None:
-            payload["input"]["steps"] = steps
+            inp["steps"] = steps
+        if size:
+            inp["size"] = size
         if self._secret:
-            payload["input"]["secret"] = self._secret
+            inp["secret"] = self._secret
+        payload: dict[str, Any] = {"input": inp}
 
         try:
             data = self._post(

{kgmodule_utils-0.4.2 → kgmodule_utils-0.4.3}/src/kg_utils/worker/ops.py

@@ -48,24 +48,39 @@ def handle_aux_ops(
             return {"error": "imagine requires a non-empty 'prompt'"}
 
         aspect = inp.get("aspect_ratio", "3:2")
+        size = inp.get("size")
         seed = inp.get("seed")
         steps = inp.get("steps")
         img_synth = image_synth_factory(inp.get("image_backend", ""))
+        seed_int = int(seed) if seed is not None else None
+        steps_int = int(steps) if steps is not None else None
 
         try:
-            b64 = img_synth.generate_b64(
-                prompt,
-                aspect_ratio=aspect,
-                seed=int(seed) if seed is not None else None,
-                steps=int(steps) if steps is not None else None,
-            )
-            return {
+            if size is not None:
+                b64 = img_synth.generate_b64(
+                    prompt,
+                    aspect_ratio=aspect,
+                    seed=seed_int,
+                    steps=steps_int,
+                    size=size,
+                )
+            else:
+                b64 = img_synth.generate_b64(
+                    prompt,
+                    aspect_ratio=aspect,
+                    seed=seed_int,
+                    steps=steps_int,
+                )
+            result = {
                 "image_b64": b64,
                 "prompt": prompt,
                 "aspect_ratio": aspect,
                 "image_model": img_synth._cfg.resolved_model(),  # pylint: disable=protected-access
                 "image_backend": img_synth._cfg.backend.value,  # pylint: disable=protected-access
             }
+            if size is not None:
+                result["size"] = size
+            return result
         except Exception as exc:  # noqa: BLE001  # pylint: disable=broad-exception-caught
             return {"error": f"image generation failed: {exc}"}