Spaces:
Running on Zero
Running on Zero
File size: 4,135 Bytes
8a28a8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """Runtime configuration shared by the web UI, MCP API, and engine."""
from __future__ import annotations
import os
from dataclasses import dataclass
def _env_bool(name: str, default: bool) -> bool:
raw = os.getenv(name)
if raw is None:
return default
return raw.strip().lower() in {"1", "true", "yes", "on"}
def _env_int(name: str, default: int, minimum: int, maximum: int) -> int:
raw = os.getenv(name)
if raw is None:
return default
try:
value = int(raw)
except (TypeError, ValueError):
return default
return max(minimum, min(maximum, value))
def _env_float(name: str, default: float, minimum: float, maximum: float) -> float:
raw = os.getenv(name)
if raw is None:
return default
try:
value = float(raw)
except (TypeError, ValueError):
return default
return max(minimum, min(maximum, value))
@dataclass(frozen=True)
class RuntimeConfig:
"""Bounded settings with safe defaults for one shared GPU process."""
# ComfyUI nodes, caches, and model state are process-global. A single
# process is therefore always serialized; scale with separate replicas.
gpu_concurrency: int = _env_int("IMAGEGEN_GPU_CONCURRENCY", 1, 1, 1)
queue_max_size: int = _env_int("IMAGEGEN_QUEUE_MAX_SIZE", 24, 1, 256)
mcp_max_pending: int = _env_int("IMAGEGEN_MCP_MAX_PENDING", 16, 1, 256)
mcp_task_retention: int = _env_int("IMAGEGEN_MCP_TASK_RETENTION", 200, 20, 5000)
max_batch_size: int = _env_int("IMAGEGEN_MAX_BATCH_SIZE", 4, 1, 16)
# ZeroGPU is best used as one bounded job at a time. The environment
# overrides are intentionally available for dedicated hardware, while the
# defaults keep the public Space predictable and affordable.
max_pk_models: int = _env_int("IMAGEGEN_MAX_PK_MODELS", 2, 2, 4)
max_multi_images: int = _env_int("IMAGEGEN_MAX_MULTI_IMAGES", 4, 1, 10)
max_plan_jobs: int = _env_int("IMAGEGEN_MAX_PLAN_JOBS", 4, 2, 12)
max_plan_outputs: int = _env_int("IMAGEGEN_MAX_PLAN_OUTPUTS", 8, 2, 16)
output_retention: int = _env_int("IMAGEGEN_OUTPUT_RETENTION", 80, 10, 1000)
max_input_megapixels: float = _env_float(
"IMAGEGEN_MAX_INPUT_MEGAPIXELS", 4.2, 0.25, 64.0
)
max_reference_megapixels: float = _env_float(
"IMAGEGEN_MAX_REFERENCE_MEGAPIXELS", 12.0, 1.0, 64.0
)
max_reference_images: int = _env_int(
"IMAGEGEN_MAX_REFERENCE_IMAGES", 10, 1, 20
)
max_output_megapixels: float = _env_float(
"IMAGEGEN_MAX_OUTPUT_MEGAPIXELS", 16.0, 1.0, 128.0
)
min_free_disk_gb: float = _env_float(
"IMAGEGEN_MIN_FREE_DISK_GB", 3.0, 0.5, 50.0
)
enable_mcp: bool = _env_bool("IMAGEGEN_ENABLE_MCP", True)
enable_startup_gpu_probe: bool = _env_bool("IMAGEGEN_STARTUP_GPU_PROBE", False)
default_language: str = os.getenv("IMAGEGEN_DEFAULT_LANGUAGE", "zh-CN")
CONFIG = RuntimeConfig()
def estimate_gpu_duration(inputs: dict) -> int:
"""Estimate a ZeroGPU reservation without asking ordinary users to guess.
This is deliberately conservative. A caller can still pass 60/90/120 via
``zero_gpu_duration``; 0 or an empty value means automatic mode.
"""
explicit = inputs.get("zero_gpu_duration")
try:
explicit_value = int(explicit or 0)
except (TypeError, ValueError):
explicit_value = 0
if explicit_value > 0:
return max(30, min(120, explicit_value))
model = str(inputs.get("model_display_name", "")).lower()
steps = max(1, int(inputs.get("num_inference_steps") or 20))
batch = max(1, int(inputs.get("batch_size") or 1))
width = max(256, int(inputs.get("width") or 1024))
height = max(256, int(inputs.get("height") or 1024))
megapixels = (width * height) / 1_000_000
fast_tokens = ("turbo", "lightning", "fast", "distilled", "schnell", "4b")
base = 35 if any(token in model for token in fast_tokens) else 60
work = steps * batch * max(0.5, megapixels)
if work >= 55 or batch >= 3 or megapixels >= 2.5:
return 120
if work >= 28:
return 90
return max(45, base)
|