horde-engine 3.4.2__tar.gz → 3.6.0__tar.gz

{horde_engine-3.4.2 → horde_engine-3.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: horde-engine
-Version: 3.4.2
+Version: 3.6.0
 Summary: A wrapper around ComfyUI to allow use by AI Horde.
 Project-URL: Homepage, https://github.com/Haidra-Org/hordelib
 Project-URL: Bug Tracker, https://github.com/Haidra-Org/hordelib/issues

{horde_engine-3.4.2 → horde_engine-3.6.0}/hordelib/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '3.4.2'
-__version_tuple__ = version_tuple = (3, 4, 2)
+__version__ = version = '3.6.0'
+__version_tuple__ = version_tuple = (3, 6, 0)
 
 __commit_id__ = commit_id = None

{horde_engine-3.4.2 → horde_engine-3.6.0}/hordelib/execution/comfy_patches.py

@@ -33,6 +33,7 @@ FORCE_LOAD_SKIP_CLASS_NAMES: tuple[str, ...] = (
     "SDXLRefiner",
     "Flux",
     "QwenImage",
+    "Lumina2",
 )
 """The comfy ``model_base`` class names hordelib must keep in lockstep with ComfyUI.
 
@@ -88,6 +89,8 @@ def _baseline_class_names() -> dict:
                 KNOWN_IMAGE_GENERATION_BASELINE.flux_schnell: ("Flux",),
                 KNOWN_IMAGE_GENERATION_BASELINE.flux_dev: ("Flux",),
                 KNOWN_IMAGE_GENERATION_BASELINE.qwen_image: ("QwenImage",),
+                # Z-Image (incl. Z-Image-Turbo) loads as comfy's Lumina2 model_base class.
+                KNOWN_IMAGE_GENERATION_BASELINE.z_image_turbo: ("Lumina2",),
             },
         )
     return BASELINE_FORCE_LOAD_CLASS_NAMES

{horde_engine-3.4.2 → horde_engine-3.6.0}/hordelib/initialisation.py

@@ -10,6 +10,7 @@ from hordelib.config_path import get_comfyui_path, get_hordelib_path, set_system
 from hordelib.consts import RELEASE_VERSION
 from hordelib.installation import EnvironmentInstaller, load_packaged_manifest
 from hordelib.utils.logger import HordeLog
+from hordelib.utils.torch_memory import AcceleratorKind
 
 _is_initialised = False
 
@@ -27,6 +28,8 @@ def initialise(
     do_not_load_model_mangers: bool = True,
     models_not_to_force_load: list[str] | None = None,
     reference_offline: bool | None = None,
+    device_index: int | None = None,
+    accelerator_kind: AcceleratorKind | None = None,
 ):
     """Initialise hordelib. This is required before using any other hordelib functions.
 
@@ -46,9 +49,32 @@ def initialise(
         reference_offline (bool | None, optional): If True, the model reference manager reads references \
             from local disk only and never downloads them (the caller/parent process owns downloading). \
             If None, defers to ``HORDE_MODEL_REFERENCE_OFFLINE``. Defaults to None.
+        device_index (int | None, optional): Pin this process to a single device by global index. \
+            When set, applies the env-var mask and any required ComfyUI CLI args via \
+            :func:`~hordelib.utils.device_pinning.device_pin_env` before torch is imported. \
+            ``None`` (default) leaves the current environment untouched (today's behavior). \
+            ``accelerator_kind`` must also be provided when this is set.
+        accelerator_kind (AcceleratorKind | None, optional): The backend kind of ``device_index``, \
+            required when ``device_index`` is set. Obtain from \
+            :func:`~hordelib.utils.torch_memory.enumerate_accelerators`. Defaults to None.
     """
     global _is_initialised
 
+    # Device pinning must happen before ComfyUI (and torch) is imported; CUDA_VISIBLE_DEVICES
+    # has no effect after torch.cuda is initialised. Apply the mask here, ahead of everything else.
+    if device_index is not None:
+        if accelerator_kind is None:
+            raise ValueError("accelerator_kind must be provided when device_index is set")
+        from hordelib.utils.device_pinning import device_pin_env
+
+        pin_env, pin_args = device_pin_env(accelerator_kind, device_index)
+        os.environ.update(pin_env)
+        if extra_comfyui_args is None:
+            extra_comfyui_args = list(pin_args)
+        else:
+            extra_comfyui_args = list(extra_comfyui_args) + pin_args
+        logger.debug("Pinned process to device index={} kind={}: env={}", device_index, accelerator_kind, pin_env)
+
     # Opt the CUDA/ROCm caching allocator into expandable segments before ComfyUI imports torch.
     # Fragmentation (a large "reserved but unallocated" pool) is a common cause of an out-of-memory
     # error even when the device still reports free memory; torch's own OOM message recommends this

{horde_engine-3.4.2 → horde_engine-3.6.0}/hordelib/model_manager/base.py

@@ -1,3 +1,4 @@
+import os
 import threading
 from abc import ABC
 from collections.abc import Callable, Iterable
@@ -14,6 +15,7 @@ from horde_model_reference import (
     file_paths_for,
     get_category_descriptor,
     horde_model_reference_paths,
+    is_present,
 )
 from horde_model_reference.model_reference_manager import ModelReferenceManager
 from horde_model_reference.model_reference_records import GenericModelRecord
@@ -539,11 +541,23 @@ class BaseModelManager[RecordT: GenericModelRecord | dict[str, Any]](ABC):
                 self.validate_model(model)
         return True
 
-    def is_model_available(self, model_name: str) -> bool:
+    def _extra_weights_roots(self) -> list[Path]:
+        """Additional weights roots to search for already-present files, from the environment.
+
+        Mirrors the worker's ``AIWORKER_EXTRA_MODEL_DIRECTORIES`` (an ``os.pathsep``-separated list of
+        weights-root directories) so a model spread across disks is judged present here exactly as the
+        worker's download planner judges it.
         """
-        :param model_name: Name of the model
-        Checks if the model is available.
-        Returns True if the model is available, False otherwise.
+        raw = os.environ.get("AIWORKER_EXTRA_MODEL_DIRECTORIES", "")
+        return [Path(entry) for entry in raw.split(os.pathsep) if entry.strip()]
+
+    def is_model_available(self, model_name: str) -> bool:
+        """Return whether *model_name*'s declared files all exist on disk (existence-only).
+
+        Presence is delegated to horde_model_reference's canonical on-disk layout, so every consumer
+        (the worker's download plan, the TUI picker, this manager) answers "is it on disk?" the same way
+        and a model placed on disk is never reported as needing download. Integrity is a separate concern:
+        a present-but-corrupt file counts as available here and is caught by ``validate_model`` (or at load).
         """
         if model_name not in self.model_reference:
             return False
@@ -551,12 +565,11 @@ class BaseModelManager[RecordT: GenericModelRecord | dict[str, Any]](ABC):
         if model_name in self.tainted_models:
             return False
 
-        model_files = self.get_model_filenames(model_name)
-        for file_entry in model_files:
-            if not self.is_file_available(file_entry["file_path"]):
-                logger.debug([file_entry["file_path"], self.is_file_available(file_entry["file_path"])])
-                return False
-        return True
+        record = self._get_generic_record(model_name)
+        if record is None:
+            return False
+
+        return is_present(record, self._weights_root, extra_roots=self._extra_weights_roots())
 
     def is_model_url_from_civitai(self, url: str) -> bool:
         return CIVITAI_API_PATH in url

{horde_engine-3.4.2 → horde_engine-3.6.0}/hordelib/pipeline/families/image.py

@@ -7,10 +7,11 @@ The predicates and priorities reproduce the legacy decision tree in
 2. cascade (remix / 2pass / base)
 3. flux
 4. qwen
-5. controlnet (annotator / hires / base)
-6. img2img-with-mask / inpainting / outpainting
-7. hires fix
-8. generic stable diffusion (also plain img2img)
+5. z-image
+6. controlnet (annotator / hires / base)
+7. img2img-with-mask / inpainting / outpainting
+8. hires fix
+9. generic stable diffusion (also plain img2img)
 
 The bindings are the typed port of the legacy ``PAYLOAD_TO_PIPELINE_PARAMETER_MAPPING``: the
 full candidate set is declared once and each template keeps the bindings whose target node
@@ -62,6 +63,15 @@ FLUX_BASELINES: frozenset[KNOWN_IMAGE_GENERATION_BASELINE] = frozenset(
     },
 )
 
+# Split-files baselines whose diffusion weights load through the bare diffusion-model loader
+# (file_type "unet") with CLIP/VAE wired from their own loader nodes, rather than a fused checkpoint.
+UNET_LOADER_BASELINES: frozenset[KNOWN_IMAGE_GENERATION_BASELINE] = frozenset(
+    {
+        KNOWN_IMAGE_GENERATION_BASELINE.qwen_image,
+        KNOWN_IMAGE_GENERATION_BASELINE.z_image_turbo,
+    },
+)
+
 
 def _has_img2img_mask(payload: ImageGenPayload) -> bool:
     if payload.source_processing != "img2img":
@@ -167,9 +177,9 @@ def _apply_model_context_step(graph: ComfyGraph, payload: ImageGenPayload, conte
         "model_loader.model_name": context.main_file,
         "model_loader.horde_model_name": context.horde_model_name,
         "model_loader.will_load_loras": context.will_load_loras,
-        # The HordeCheckpointLoader needs to know what file to load; "unet" routes qwen
-        # through the diffusion-model loader, None keeps normal SD checkpoints working.
-        "model_loader.file_type": "unet" if context.baseline is KNOWN_IMAGE_GENERATION_BASELINE.qwen_image else None,
+        # The HordeCheckpointLoader needs to know what file to load; "unet" routes the split-files
+        # baselines (qwen, z-image) through the diffusion-model loader, None keeps normal SD checkpoints working.
+        "model_loader.file_type": "unet" if context.baseline in UNET_LOADER_BASELINES else None,
         "model_loader_stage_c.ckpt_name": context.extra_files.get("stable_cascade_stage_c"),
         "model_loader_stage_c.model_name": context.extra_files.get("stable_cascade_stage_c"),
         "model_loader_stage_c.horde_model_name": context.horde_model_name,
@@ -376,6 +386,11 @@ def build_default_registry() -> PipelineRegistry[ImageGenPayload, ModelContext]:
             predicate=lambda p, c: c.baseline is KNOWN_IMAGE_GENERATION_BASELINE.qwen_image,
             priority=86,
         ),
+        PipelineSpec(
+            template=_template("z_image"),
+            predicate=lambda p, c: c.baseline is KNOWN_IMAGE_GENERATION_BASELINE.z_image_turbo,
+            priority=85,
+        ),
         # ControlNet
         PipelineSpec(
             template=_template("controlnet_annotator"),

horde_engine-3.6.0/hordelib/pipelines/pipeline_z_image.json

@@ -0,0 +1,128 @@
+{
+  "3": {
+    "inputs": {
+      "seed": 556312751097931,
+      "steps": 9,
+      "cfg": 1,
+      "sampler_name": "euler",
+      "scheduler": "simple",
+      "denoise": 1,
+      "model": [
+        "37",
+        0
+      ],
+      "positive": [
+        "6",
+        0
+      ],
+      "negative": [
+        "7",
+        0
+      ],
+      "latent_image": [
+        "58",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "sampler"
+    }
+  },
+  "6": {
+    "inputs": {
+      "text": "a photograph of a horse",
+      "clip": [
+        "38",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "prompt"
+    }
+  },
+  "7": {
+    "inputs": {
+      "text": "",
+      "clip": [
+        "38",
+        0
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "negative_prompt"
+    }
+  },
+  "8": {
+    "inputs": {
+      "samples": [
+        "3",
+        0
+      ],
+      "vae": [
+        "39",
+        0
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "vae_decode"
+    }
+  },
+  "37": {
+    "inputs": {
+      "unet_name": "z_image_turbo_bf16.safetensors",
+      "weight_dtype": "default"
+    },
+    "class_type": "UNETLoader",
+    "_meta": {
+      "title": "model_loader"
+    }
+  },
+  "38": {
+    "inputs": {
+      "clip_name": "qwen_3_4b.safetensors",
+      "type": "lumina2",
+      "device": "default"
+    },
+    "class_type": "CLIPLoader",
+    "_meta": {
+      "title": "clip_loader"
+    }
+  },
+  "39": {
+    "inputs": {
+      "vae_name": "ae.safetensors"
+    },
+    "class_type": "VAELoader",
+    "_meta": {
+      "title": "vae_loader"
+    }
+  },
+  "58": {
+    "inputs": {
+      "width": 1024,
+      "height": 1024,
+      "batch_size": 1
+    },
+    "class_type": "EmptySD3LatentImage",
+    "_meta": {
+      "title": "empty_latent_image"
+    }
+  },
+  "75": {
+    "inputs": {
+      "filename_prefix": "ComfyUI",
+      "images": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "SaveImage",
+    "_meta": {
+      "title": "output_image"
+    }
+  }
+}

horde_engine-3.6.0/hordelib/utils/device_pinning.py

@@ -0,0 +1,54 @@
+"""Single source of truth for masking a subprocess to one compute device.
+
+Each inference child process is pinned to exactly one device so that every
+single-device assumption in ComfyUI and hordelib (``cuda:0``, module-level
+``total_vram``, ``UserSettings._vram_to_leave_free_mb``) stays correct without
+modification inside that child. The orchestrator (which is torch-free) calls
+:func:`device_pin_env` to get the env-var patch and extra ComfyUI CLI args it
+must apply before spawning the child; the child itself applies them before
+calling ``hordelib.initialise()``.
+
+:func:`device_pin_env` is the only function here. All callers (worker entry
+points, ``hordelib.initialise``'s optional convenience path) must go through
+it so that backend-specific masking details stay in one place.
+"""
+
+from __future__ import annotations
+
+from hordelib.utils.torch_memory import AcceleratorKind
+
+
+def device_pin_env(kind: AcceleratorKind, index: int) -> tuple[dict[str, str], list[str]]:
+    """Return the env-var patch and extra ComfyUI args that restrict a process to device ``index``.
+
+    After applying both, the OS/driver presents only that one device to the process and it
+    becomes ``cuda:0`` (or the backend equivalent), so all single-device assumptions inside
+    hordelib and the vendored ComfyUI remain correct without any code changes there.
+
+    Args:
+        kind: The :class:`~hordelib.utils.torch_memory.AcceleratorKind` of the target device.
+        index: The **global** device index (as reported by
+            :func:`~hordelib.utils.torch_memory.enumerate_accelerators`).
+
+    Returns:
+        A 2-tuple ``(env_vars, extra_args)`` where ``env_vars`` is a dict to merge into
+        ``os.environ`` and ``extra_args`` is a list of CLI arguments to append to
+        ``extra_comfyui_args``. Either or both may be empty when the backend needs no masking
+        (cpu, mps).
+    """
+    if kind is AcceleratorKind.cuda:
+        return {"CUDA_VISIBLE_DEVICES": str(index)}, []
+
+    if kind is AcceleratorKind.rocm:
+        # ROCm presents through torch.cuda, so CUDA_VISIBLE_DEVICES is honoured too; keep them in sync.
+        return {"HIP_VISIBLE_DEVICES": str(index), "CUDA_VISIBLE_DEVICES": str(index)}, []
+
+    if kind is AcceleratorKind.xpu:
+        return {"ZE_AFFINITY_MASK": str(index)}, []
+
+    if kind is AcceleratorKind.directml:
+        # DirectML device selection happens via a ComfyUI CLI flag, not an env var.
+        return {}, ["--directml", str(index)]
+
+    # cpu and mps have no multi-device isolation mechanism; masking is a no-op.
+    return {}, []

{horde_engine-3.4.2 → horde_engine-3.6.0}/hordelib/utils/gpuinfo.py

@@ -25,15 +25,25 @@ _AVERAGE_WINDOW_SECONDS = 60 * 5
 class GPUInfo:
     """Collects per-device GPU statistics, keeping rolling averages of load, temperature, and power."""
 
-    def __init__(self) -> None:
-        """Initialize the collector for the device named by ``CUDA_VISIBLE_DEVICES`` (default 0)."""
+    def __init__(self, device_index: int | None = None) -> None:
+        """Initialize the collector for the given device index.
+
+        Args:
+            device_index: The global NVML/CUDA device index to query. When ``None`` (default), reads
+                ``CUDA_VISIBLE_DEVICES`` (which is ``0`` inside a pinned subprocess, and the real
+                device index in the unmasked orchestrator). Pass an explicit index when an unmasked
+                caller needs to survey a specific card without masking the process.
+        """
         self.avg_load: list[int] = []
         self.avg_temp: list[int] = []
         self.avg_power: list[int] = []
         self._average_window_samples = _AVERAGE_SAMPLES_PER_SECOND * _AVERAGE_WINDOW_SECONDS
         if not self.is_nvidia():
             logger.warning("Detailed GPU info (load/temp/power) is only available on NVIDIA GPUs")
-        self.device = int(os.getenv("CUDA_VISIBLE_DEVICES", "0"))
+        if device_index is not None:
+            self.device = device_index
+        else:
+            self.device = int(os.getenv("CUDA_VISIBLE_DEVICES", "0"))
 
     def is_nvidia(self) -> bool:
         """Return whether the active torch build targets NVIDIA CUDA (not ROCm/HIP)."""

{horde_engine-3.4.2 → horde_engine-3.6.0}/pyproject.toml

@@ -298,6 +298,7 @@ markers = [
     "default_sd15_model: marks tests as default_sd15_model (deselect with '-m \"not default_sd15_model\"')",
     "default_sdxl_model: marks tests as default_sdxl_model (deselect with '-m \"not default_sdxl_model\"')",
     "default_qwen_model: marks tests as default_qwen_model (deselect with '-m \"not default_qwen_model\"')",
+    "default_z_image_turbo_model: marks tests as default_z_image_turbo_model (deselect with '-m \"not default_z_image_turbo_model\"')",
     "refined_sdxl_model: marks tests as refined_sdxl_model (deselect with '-m \"not refined_sdxl_model\"')",
     "patch_harness: instrumentation-only tests for ComfyUI monkeypatch impact",
 ]