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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kgmodule-utils
-Version: 0.4.2
+Version: 0.4.4
 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.4}/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "kgmodule-utils"
-version = "0.4.2"
+version = "0.4.4"
 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.4}/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.4"

{kgmodule_utils-0.4.2 → kgmodule_utils-0.4.4}/src/kg_utils/embedder.py

@@ -34,12 +34,12 @@ License: Elastic 2.0
 
 from __future__ import annotations
 
+import importlib
 import os
 from typing import Any
 
 from kg_utils.embed import DEFAULT_MODEL, KNOWN_MODELS, resolve_model_path
 
-
 # ---------------------------------------------------------------------------
 # Abstract base
 # ---------------------------------------------------------------------------
@@ -75,7 +75,7 @@ class Embedder:
 # ---------------------------------------------------------------------------
 
 
-def load_sentence_transformer(model_name: str = DEFAULT_MODEL) -> Any:
+def load_sentence_transformer(model_name: str = DEFAULT_MODEL, device: str | None = None) -> Any:
     """Load a ``SentenceTransformer`` with the canonical safe-load sequence.
 
     Resolution order:
@@ -90,13 +90,22 @@ def load_sentence_transformer(model_name: str = DEFAULT_MODEL) -> Any:
     retry loops leave stale thread state that causes SIGBUS on the first
     ``encode()`` call.
 
+    Device precedence: explicit *device* arg > ``KG_EMBED_DEVICE`` env >
+    auto-detect.  The env var exists because ``spawn``-based embedding workers
+    inherit ``os.environ`` but can't easily receive a Python arg — without a way
+    to pin the device, each worker auto-selects MPS and N parallel workers stack
+    N GPU allocations into an OOM.  So CPU multiprocessing embedding on Apple
+    Silicon is only safe with this knob.
+
     :param model_name: HuggingFace model ID or KNOWN_MODELS alias.
+    :param device: Explicit device (``"cpu"``/``"mps"``/``"cuda"``).  ``None``
+        falls back to ``KG_EMBED_DEVICE`` then CUDA→MPS→CPU auto-detect.
     :return: Loaded ``SentenceTransformer`` instance.
     """
-    from sentence_transformers import SentenceTransformer  # pylint: disable=import-outside-toplevel
+    SentenceTransformer = importlib.import_module("sentence_transformers").SentenceTransformer
 
     try:
-        from transformers import logging as hf_logging  # pylint: disable=import-outside-toplevel
+        hf_logging = importlib.import_module("transformers.logging")
 
         hf_logging.set_verbosity_error()
         # TQDM_DISABLE alone misses transformers' _tqdm_active gate
@@ -106,9 +115,18 @@ def load_sentence_transformer(model_name: str = DEFAULT_MODEL) -> Any:
 
     os.environ["TQDM_DISABLE"] = "1"
 
-    import torch  # pylint: disable=import-outside-toplevel
-
-    if torch.cuda.is_available():
+    torch = importlib.import_module("torch")
+
+    # Device precedence: explicit arg > KG_EMBED_DEVICE env > auto-detect.
+    # The env var lets spawn-based embedding workers (which inherit os.environ
+    # but can't easily receive a Python arg) pin to e.g. CPU — without it each
+    # worker auto-selects MPS and N parallel workers stack N GPU allocations →
+    # MPS OOM. The override is why CPU multiprocessing embedding is safe on
+    # Apple Silicon.
+    sel = (device or os.environ.get("KG_EMBED_DEVICE", "")).strip().lower()
+    if sel:
+        device = sel
+    elif torch.cuda.is_available():
         device = "cuda"
     else:
         try:
@@ -158,7 +176,7 @@ class SentenceTransformerEmbedder(Embedder):
 
     def __init__(self, model_name: str = DEFAULT_MODEL) -> None:
         try:
-            from transformers import logging as hf_logging  # pylint: disable=import-outside-toplevel
+            hf_logging = importlib.import_module("transformers.logging")
 
             hf_logging.set_verbosity_error()
             hf_logging.disable_progress_bar()
@@ -188,7 +206,7 @@ class SentenceTransformerEmbedder(Embedder):
         :param texts: Input strings.
         :param encode_batch_size: Passed to ``model.encode()`` — tune down if OOM on MPS.
         """
-        import numpy as np  # pylint: disable=import-outside-toplevel
+        np = importlib.import_module("numpy")
 
         vecs = self.model.encode(
             texts,
@@ -200,7 +218,7 @@ class SentenceTransformerEmbedder(Embedder):
 
     def embed_query(self, query: str) -> list[float]:
         """Embed a single query string into a float32 vector."""
-        import numpy as np  # pylint: disable=import-outside-toplevel
+        np = importlib.import_module("numpy")
 
         vec = self.model.encode([query], normalize_embeddings=True)[0]
         return list(np.asarray(vec, dtype="float32").tolist())
@@ -233,7 +251,7 @@ def wrap_embedder(st_model: Any, model_name: str = DEFAULT_MODEL) -> Embedder:
     :param model_name: Model name stored as metadata on the wrapper.
     :return: An :class:`Embedder` that delegates all calls to *st_model*.
     """
-    import numpy as np  # pylint: disable=import-outside-toplevel
+    np = importlib.import_module("numpy")
 
     resolved = KNOWN_MODELS.get(model_name, model_name)
     _dim_fn = getattr(st_model, "get_embedding_dimension", None) or getattr(

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

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import base64
+import importlib
 import random
 from io import BytesIO
 from typing import TYPE_CHECKING, Any
@@ -51,6 +52,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
 # ---------------------------------------------------------------------------
@@ -79,7 +103,10 @@ class ImageSynthesizer:
     def _load_mflux(self, model_name: str) -> Any:
         if self._mflux_model is not None and self._mflux_model_name == model_name:
             return self._mflux_model
-        from mflux.models.flux2.variants.txt2img.flux2_klein import Flux2Klein  # type: ignore
+
+        Flux2Klein = importlib.import_module(
+            "mflux.models.flux2.variants.txt2img.flux2_klein"
+        ).Flux2Klein
 
         self._mflux_model = Flux2Klein(model_path=model_name)
         self._mflux_model_name = model_name
@@ -97,6 +124,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 +133,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 +149,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 +158,7 @@ class ImageSynthesizer:
                 seed=seed,
                 steps=effective_steps,
                 aspect_ratio=aspect_ratio,
+                size=size,
             )
         else:
             return self._generate_openai(
@@ -142,6 +175,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 +186,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 +208,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 +232,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.4}/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.4}/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}"}