horde-engine 3.3.0__tar.gz → 3.3.2__tar.gz

{horde_engine-3.3.0 → horde_engine-3.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: horde-engine
-Version: 3.3.0
+Version: 3.3.2
 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.3.0 → horde_engine-3.3.2}/hordelib/_version.py

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

{horde_engine-3.3.0 → horde_engine-3.3.2}/hordelib/api.py

@@ -23,6 +23,7 @@ from hordelib.execution.sampling_lease import (
 )
 from hordelib.feature_impact import (
     FEATURE_KIND,
+    FEATURE_PHASE,
     BaselineBurden,
     BurdenEstimate,
     CalibrationSample,
@@ -98,6 +99,7 @@ if TYPE_CHECKING:
 __all__ = [
     "CONTROLNET_ANNOTATOR_DOWNLOAD_BYTES",
     "FEATURE_KIND",
+    "FEATURE_PHASE",
     "AcceleratorInfo",
     "AcceleratorKind",
     "BaselineBurden",

{horde_engine-3.3.0 → horde_engine-3.3.2}/hordelib/feature_impact.py

@@ -37,6 +37,20 @@ class FEATURE_KIND(StrEnum):
     strip_background = auto()
 
 
+class FEATURE_PHASE(StrEnum):
+    """When in a job's lifecycle a feature's VRAM is consumed.
+
+    A consumer that schedules several jobs on one device needs this distinction: the ``sampling`` cost
+    is co-resident with the checkpoint while it denoises, whereas a ``post_processing`` cost (an upscaler
+    or face-fixer) is claimed *after* sampling completes and the checkpoint may already have been evicted.
+    Summing both into a single per-job figure hides that the post-processing peaks of distinct in-flight
+    jobs can align in time and over-commit the card, so the estimate is reported per phase.
+    """
+
+    sampling = auto()
+    post_processing = auto()
+
+
 class DownloadTrigger(BaseModel):
     """Whether using a feature/model can trigger an ad-hoc network download."""
 
@@ -54,6 +68,8 @@ class FeatureImpact(BaseModel):
     """Activation-scaling component, applied per output megapixel."""
     ram_delta_mb: int
     download: DownloadTrigger
+    phase: FEATURE_PHASE = FEATURE_PHASE.sampling
+    """The job phase this feature's VRAM is consumed in; drives the per-phase split of the estimate."""
     applies_to_baselines: list[str] | None = None
     """Baseline values this feature is available for (None = all)."""
     notes: str = ""
@@ -91,6 +107,17 @@ class BurdenEstimate(BaseModel):
     """The estimated total burden of one job configuration."""
 
     vram_mb: int
+    """Total peak VRAM (``vram_sampling_mb + vram_post_processing_mb``); kept for back-compat callers."""
+    vram_sampling_mb: int = 0
+    """VRAM resident during sampling: the baseline plus every sampling-phase feature delta.
+
+    Defaults to 0 so a ``BurdenEstimate`` deserialized from an older producer is still constructible; live
+    estimates from :func:`estimate_job_burden` always populate it."""
+    vram_post_processing_mb: int = 0
+    """Marginal VRAM peak of the post-processing phase (upscalers/face-fixers), claimed after sampling.
+
+    This is the figure a scheduler reserves against concurrent dispatch, since it lands while the slot has
+    already been released for the next job. Zero when the job has no post-processing features."""
     ram_mb: int
     disk_bytes_needed: int
     downloads_expected: list[DownloadTrigger]
@@ -299,6 +326,7 @@ _FEATURE_SEEDS: list[FeatureImpact] = [
         vram_delta_per_megapixel_mb=300,
         ram_delta_mb=2000,
         download=DownloadTrigger(triggered=True, typical_size_mb=70),
+        phase=FEATURE_PHASE.post_processing,
         notes="ESRGAN-family upscalers; model is small but activations on large outputs are not.",
     ),
     FeatureImpact(
@@ -306,6 +334,7 @@ _FEATURE_SEEDS: list[FeatureImpact] = [
         vram_delta_mb=1500,
         ram_delta_mb=2000,
         download=DownloadTrigger(triggered=True, typical_size_mb=350),
+        phase=FEATURE_PHASE.post_processing,
         notes="GFPGAN/CodeFormer weights.",
     ),
     FeatureImpact(
@@ -313,6 +342,7 @@ _FEATURE_SEEDS: list[FeatureImpact] = [
         vram_delta_mb=1000,
         ram_delta_mb=1500,
         download=DownloadTrigger(triggered=True, typical_size_mb=170),
+        phase=FEATURE_PHASE.post_processing,
     ),
 ]
 
@@ -340,9 +370,28 @@ def estimate_job_burden(
     batch: int = 1,
     features: list[FEATURE_KIND] | None = None,
     model_record: ImageGenerationModelRecord | None = None,
+    aux_model_weights_mb: dict[FEATURE_KIND, float] | None = None,
+    post_processing_upscale_factor: float = 1.0,
 ) -> BurdenEstimate:
     """Estimate the total resource burden of one job configuration.
 
+    The VRAM total is reported both as a single ``vram_mb`` (back-compat) and split by
+    :class:`FEATURE_PHASE`: ``vram_sampling_mb`` (baseline plus sampling-phase features) and
+    ``vram_post_processing_mb`` (the marginal upscaler/face-fixer peak that lands after sampling). A
+    scheduler reserves the post-processing figure against concurrent dispatch because it is claimed once
+    the inference slot has already been released for the next job.
+
+    ``aux_model_weights_mb`` lets a caller that has resolved a specific auxiliary model (e.g. a particular
+    ESRGAN upscaler or controlnet) supply its real resident weight (MB), replacing that feature's flat
+    ``vram_delta_mb`` term while keeping the activation (per-megapixel) term. Unsupplied features keep the
+    conservative seed.
+
+    ``post_processing_upscale_factor`` is the linear scale of the requested upscaler (see
+    :func:`hordelib.pipeline.constants.max_upscale_factor`). The upscale activation peak scales with the
+    *output* megapixels, so the ``post_processing_upscale`` feature's per-megapixel term is applied against
+    ``factor**2`` the generation megapixels: a 4x upscale of a 1 MP image works a 16 MP tensor, the dominant
+    cost of the post-processing phase. The default of 1.0 leaves the activation at generation resolution.
+
     Never raises on unknown baselines: pre-flight callers need an answer for every
     job, so unknown baselines use a heavy fallback seed and are flagged via
     ``baseline_known=False``.
@@ -354,7 +403,8 @@ def estimate_job_burden(
 
     megapixels = (width * height) / 1_000_000
 
-    vram_mb = burden.vram_base_mb + round(burden.vram_per_megapixel_mb * megapixels * batch)
+    vram_sampling_mb = burden.vram_base_mb + round(burden.vram_per_megapixel_mb * megapixels * batch)
+    vram_post_processing_mb = 0
     ram_mb = burden.ram_base_mb
     downloads_expected: list[DownloadTrigger] = []
 
@@ -364,7 +414,20 @@ def estimate_job_burden(
             continue
         if impact.applies_to_baselines is not None and baseline not in impact.applies_to_baselines:
             continue
-        vram_mb += impact.vram_delta_mb + round(impact.vram_delta_per_megapixel_mb * megapixels)
+        if aux_model_weights_mb is not None and kind in aux_model_weights_mb:
+            weight_mb = round(aux_model_weights_mb[kind])
+        else:
+            weight_mb = impact.vram_delta_mb
+        # The upscaler enlarges the image, so its activation peak scales with the *output* megapixels
+        # (factor**2 the generation megapixels); every other feature works at generation resolution.
+        activation_megapixels = megapixels
+        if kind == FEATURE_KIND.post_processing_upscale and post_processing_upscale_factor > 1.0:
+            activation_megapixels = megapixels * (post_processing_upscale_factor**2)
+        vram_delta = weight_mb + round(impact.vram_delta_per_megapixel_mb * activation_megapixels)
+        if impact.phase == FEATURE_PHASE.post_processing:
+            vram_post_processing_mb += vram_delta
+        else:
+            vram_sampling_mb += vram_delta
         ram_mb += impact.ram_delta_mb
         if impact.download.triggered:
             downloads_expected.append(impact.download)
@@ -376,7 +439,9 @@ def estimate_job_burden(
     disk_bytes += sum((trigger.typical_size_mb or 0) * 1024 * 1024 for trigger in downloads_expected)
 
     return BurdenEstimate(
-        vram_mb=vram_mb,
+        vram_mb=vram_sampling_mb + vram_post_processing_mb,
+        vram_sampling_mb=vram_sampling_mb,
+        vram_post_processing_mb=vram_post_processing_mb,
         ram_mb=ram_mb,
         disk_bytes_needed=disk_bytes,
         downloads_expected=downloads_expected,

{horde_engine-3.3.0 → horde_engine-3.3.2}/hordelib/pipeline/constants.py

@@ -93,6 +93,47 @@ def controlnet_annotator_download_bytes(control_types: Iterable[str | None]) ->
     return sum(CONTROLNET_ANNOTATOR_DOWNLOAD_BYTES.get(control_type, 0) for control_type in distinct)
 
 
+# Horde upscaler name (KNOWN_UPSCALERS value) to its linear scale factor. The upscaler enlarges the
+# generated image by this factor on each axis, so the post-processing activation peak scales with the
+# *output* megapixels (factor**2 the generation megapixels), not the generation resolution. A 4x upscale
+# of a 1 MP image produces a 16 MP tensor, which is the dominant VRAM cost of the post-processing phase
+# (the model weights themselves are tens of MB). Used by the feature-impact estimate to size that peak.
+UPSCALER_SCALE_FACTORS = {
+    "BACKEND_DEFAULT": 4,  # the worker's default upscaler is a 4x ESRGAN; assume the larger factor
+    "RealESRGAN_x4plus": 4,
+    "RealESRGAN_x2plus": 2,
+    "RealESRGAN_x4plus_anime_6B": 4,
+    "NMKD_Siax": 4,
+    "4x_AnimeSharp": 4,
+}
+"""Horde upscaler name to its linear (per-axis) scale factor (ROM)."""
+
+_DEFAULT_UPSCALE_FACTOR = 4
+"""Assumed factor for an upscaler absent from the ROM: err high so the activation peak is not under-sized."""
+
+
+def upscaler_scale_factor(name: str | None) -> int:
+    """Return the linear scale factor for upscaler *name*, or the conservative default when unknown.
+
+    ``None`` returns 1 (no upscaler, no enlargement). An unrecognised upscaler returns
+    :data:`_DEFAULT_UPSCALE_FACTOR` rather than 1, so a new upscaler the ROM has not learned yet
+    over-reserves rather than under-reserving the post-processing activation peak.
+    """
+    if name is None:
+        return 1
+    return UPSCALER_SCALE_FACTORS.get(name, _DEFAULT_UPSCALE_FACTOR)
+
+
+def max_upscale_factor(names: Iterable[str | None]) -> int:
+    """Return the largest linear scale factor among *names*, or 1 when none enlarge the image.
+
+    A job may request several upscalers; the output size (and thus the activation peak) is driven by the
+    largest factor. ``None`` entries and an empty iterable contribute the no-op factor of 1.
+    """
+    factors = [upscaler_scale_factor(name) for name in names if name is not None]
+    return max(factors) if factors else 1
+
+
 SOURCE_IMAGE_PROCESSING_OPTIONS = ["img2img", "inpainting", "outpainting", "remix"]
 
 SCHEDULERS = ["normal", "karras", "simple", "ddim_uniform", "sgm_uniform", "exponential"]