inline-core 1.1.1__py3-none-any.whl

inline_core/models/zimage/runner.py

@@ -0,0 +1,386 @@
+"""Z-Image Turbo runner: prompt (+ optional image) -> one rendered take.
+
+A single generation node, ``alibaba/z-image-turbo``, backed by diffusers' ``ZImagePipeline``
+(text-to-image) and ``ZImageImg2ImgPipeline`` (when an image is wired in). The heavy pipeline is
+built once and cached across runs; only the descriptor is cheap. Placement (device, dtype, offload,
+tiling) comes from the DevicePolicy — the runner never picks a device itself. Decoded images are
+handed to the TakeStore, which owns bytes/hash/uri.
+
+torch + diffusers are imported at module top on purpose: an absent ``zimage`` extra makes this
+import raise, and ``server.bootstrap`` skips the model (best-effort) so the engine still boots.
+"""
+
+from __future__ import annotations
+
+import os
+import random
+from threading import Lock
+from typing import Any
+
+import torch
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+    ZImageImg2ImgPipeline,
+    ZImagePipeline,
+    ZImageTransformer2DModel,
+)
+
+from ...device.policy import DevicePolicy, Placement, Profile
+from ...device.types import DeviceKind, DType
+from ...errors import CancelledError, ComponentError
+from ...graph.descriptor import NodeDescriptor, ParamField, Port, Widget
+from ...graph.runners import NodeResult, NodeRunner
+from ...graph.schema import Node, PortKind
+from ...media import MediaKind
+from ...runtime.context import ExecutionContext
+from ...runtime.progress import Phase, ProgressEvent
+from ...runtime.store import TakeStore
+from ...takes import AssetRef
+from . import requirements as reqs
+
+# The models this node needs — the diffusion transformer plus a VAE, text-encoder, tokenizer and
+# scheduler — are assembled entirely from files under models/ (see `requirements.py`). **Nothing is
+# ever downloaded here.** Every diffusers/transformers load below runs with local_files_only=True,
+# so a missing model is a clear error pointing at the node's model popup — never a silent fetch from
+# Hugging Face. Models arrive by exactly two paths: the user drops files under models/, or the popup
+# downloads them into models/.
+_SEED_MAX = 2**31 - 1
+
+
+ZIMAGE = NodeDescriptor(
+    type="alibaba/z-image-turbo",
+    title="Z-Image Turbo",
+    category="Generate",
+    icon="wand",
+    output_kind=MediaKind.IMAGE,
+    inputs=(
+        Port("prompt", "Prompt", PortKind.TEXT, required=True),
+        # Optional: wire an image to run img2img instead of text-to-image.
+        Port("image", "Image (img2img)", PortKind.IMAGE, required=False),
+    ),
+    outputs=(Port("image", "Image", PortKind.IMAGE),),
+    params=(
+        ParamField("negative_prompt", "Negative prompt", Widget.TEXTAREA, ""),
+        ParamField("width", "Width", Widget.NUMBER, 1024, min=256, max=2048, step=64),
+        ParamField("height", "Height", Widget.NUMBER, 1024, min=256, max=2048, step=64),
+        # Z-Image-Turbo is distilled: ~8 steps, CFG off (guidance 0). See the model card.
+        ParamField("steps", "Steps", Widget.NUMBER, 8, min=1, max=100, step=1),
+        ParamField("guidance", "Guidance (CFG)", Widget.NUMBER, 0.0, min=0.0, max=20.0, step=0.5),
+        # img2img only: how far to move from the input image (0 = keep, 1 = ignore).
+        ParamField(
+            "strength", "Denoise strength", Widget.NUMBER, 0.6, min=0.0, max=1.0, step=0.05,
+            advanced=True,
+        ),
+        ParamField("seed", "Seed (-1 = random)", Widget.SEED, -1),
+        # Advanced, optional: pick a specific diffusion file. "" = auto (the single file found under
+        # models/diffusion_models/). Lives behind the Adjust panel so the node stays one-click.
+        ParamField(
+            "model", "Model file (auto)", Widget.SELECT, "",
+            options_from="diffusion_models", advanced=True,
+        ),
+    ),
+)
+
+
+def register_zimage(registry: Any, store: TakeStore, policy: DevicePolicy) -> None:
+    """Register the Z-Image node and its runner. Called best-effort by server.bootstrap."""
+    registry.register(ZIMAGE, ZImageRunner(store, policy))
+
+
+class ZImageRunner(NodeRunner):
+    produces_takes = True
+
+    def __init__(self, store: TakeStore, policy: DevicePolicy) -> None:
+        self._store = store
+        self._policy = policy
+
+    def run(self, node: Node, inputs: dict[str, list[Any]], ctx: ExecutionContext) -> NodeResult:
+        prompt = _first_str(inputs.get("prompt"))
+        if not prompt:
+            raise ComponentError("Z-Image needs a prompt.")
+        params = {**ZIMAGE.defaults(), **node.params}
+        width, height = int(params["width"]), int(params["height"])
+        steps = max(1, int(params["steps"]))
+        guidance = float(params["guidance"])
+        negative = str(params.get("negative_prompt") or "").strip() or None
+        seed = _resolve_seed(params.get("seed"))
+        image_ref = _first(inputs.get("image"))
+        img2img = image_ref is not None
+
+        # No hidden downloads: if a required model isn't on disk, fail fast with a message that
+        # points at the node's model popup instead of letting diffusers silently fetch it.
+        missing = [c.label for c in reqs.zimage_requirements(params) if not c.present]
+        if missing:
+            raise ComponentError(
+                "Z-Image models missing: "
+                + ", ".join(missing)
+                + ". Download them from the node's model popup (the hint on the node)."
+            )
+        resolved = reqs.resolve_diffusion(params)
+        if resolved is None:  # defensive: the missing-check above already covers this
+            raise ComponentError("Z-Image diffusion model not found in models/diffusion_models/.")
+        source, mode = resolved
+
+        ctx.emitter.emit(_progress(ctx, node, Phase.LOADING, 0.0, status="Loading model…"))
+        pipe = _load_pipeline(self._policy, img2img=img2img, source=source, mode=mode)
+
+        placement = self._policy.placement("denoiser")
+        gen_device = "cpu" if (placement.offload or self._policy.profile is Profile.CPU) else str(
+            placement.device
+        )
+        generator = torch.Generator(device=gen_device).manual_seed(seed)
+
+        def on_step_end(_pipe: Any, step: int, _t: Any, kwargs: dict[str, Any]) -> dict[str, Any]:
+            if ctx.cancel.cancelled:
+                raise CancelledError("Run cancelled.")
+            done = step + 1
+            ctx.emitter.emit(
+                _progress(
+                    ctx,
+                    node,
+                    Phase.SAMPLE,
+                    done / steps,
+                    step=done,
+                    step_count=steps,
+                    status=f"Step {done}/{steps}",
+                )
+            )
+            return kwargs
+
+        call: dict[str, Any] = dict(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=guidance,
+            generator=generator,
+            output_type="pil",
+            callback_on_step_end=on_step_end,
+        )
+        if negative is not None:
+            call["negative_prompt"] = negative
+        if img2img:
+            call["image"] = _load_image(image_ref)
+            call["strength"] = float(params.get("strength", 0.6))
+
+        image = pipe(**call).images[0]
+
+        ctx.emitter.emit(_progress(ctx, node, Phase.SAVE, 1.0, status="Saving…"))
+        take = self._store.save(
+            ctx.run_id,
+            node.id,
+            image,
+            {
+                "model": source,
+                "prompt": prompt,
+                "negative_prompt": negative or "",
+                "width": width,
+                "height": height,
+                "steps": steps,
+                "guidance": guidance,
+                "seed": seed,
+                **({"strength": call["strength"]} if img2img else {}),
+            },
+        )
+        return NodeResult(outputs={"image": take}, takes=[take])
+
+
+# --- pipeline cache -----------------------------------------------------------------------------
+
+# Keyed by (model source, img2img). Built once; diffusers pipelines are not thread-safe, but the run
+# manager executes one run at a time (workers=1). The lock guards concurrent first-time builds.
+_PIPELINES: dict[tuple[str, bool], Any] = {}
+_LOCK = Lock()
+
+
+def _load_pipeline(policy: DevicePolicy, *, img2img: bool, source: str, mode: str) -> Any:
+    key = (source, img2img)
+    with _LOCK:
+        cached = _PIPELINES.get(key)
+        if cached is not None:
+            return cached
+        # An img2img pipe can reuse the base pipe's already-placed weights (no second load).
+        base = _PIPELINES.get((source, False))
+        if img2img and base is not None:
+            pipe = ZImageImg2ImgPipeline.from_pipe(base)
+        else:
+            dtype = _torch_dtype(policy.placement("denoiser"))
+            pipe = _build_pipeline(source, mode=mode, img2img=img2img, dtype=dtype)
+            _configure(pipe, policy)
+        _PIPELINES[key] = pipe
+        return pipe
+
+
+def _build_pipeline(source: str, *, mode: str, img2img: bool, dtype: Any) -> Any:
+    """Build a Z-Image pipeline **offline** — never touching the network.
+
+    Two shapes, both resolved from files under ``models/`` (see ``requirements.py``):
+      - ``mode == "pipeline"``: ``source`` is a whole diffusers folder (``model_index.json`` + all
+        components). Loaded with ``local_files_only=True``.
+      - ``mode == "single_file"``: ``source`` is a lone diffusion transformer file. We load the
+        transformer from it and assemble the pipeline from a **local** VAE + text-encoder +
+        tokenizer (present is guaranteed by the missing-check) plus a default flow-match scheduler.
+
+    Drop-in / override: a VAE in ``vae/`` (a ``.safetensors`` or a diffusers dir), an HF-format
+    text-encoder dir in ``text_encoders/`` (a bare weights file can't carry its config), or point
+    ``INLINE_ZIMAGE_VAE`` / ``INLINE_ZIMAGE_TEXT_ENCODER`` at them. Nothing is fetched — get missing
+    pieces via the node's model popup, which writes them under ``models/``.
+    """
+    cls = ZImageImg2ImgPipeline if img2img else ZImagePipeline
+    if mode == "pipeline":
+        return cls.from_pretrained(source, torch_dtype=dtype, local_files_only=True)
+
+    transformer = ZImageTransformer2DModel.from_single_file(
+        source, torch_dtype=dtype, local_files_only=True
+    )
+    vae = _load_local_vae(dtype)
+    text = _load_local_text_encoder(dtype)
+    if vae is None or text is None:
+        raise ComponentError(
+            "Z-Image needs a local VAE and text-encoder for a single-file diffusion model. "
+            "Download them from the node's model popup."
+        )
+    text_encoder, tokenizer = text
+    scheduler = _load_scheduler()
+    return cls(
+        scheduler=scheduler,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        transformer=transformer,
+    )
+
+
+def _load_scheduler() -> Any:
+    """The flow-match scheduler: from a local pipeline folder's ``scheduler/`` if one exists, else
+    the library default. Config-only, so this never downloads."""
+    pipe_dir = reqs.pipeline_dir(reqs.diffusion_root())
+    if pipe_dir is not None:
+        try:
+            return FlowMatchEulerDiscreteScheduler.from_pretrained(
+                str(pipe_dir), subfolder="scheduler", local_files_only=True
+            )
+        except (OSError, ValueError):
+            pass
+    return FlowMatchEulerDiscreteScheduler()
+
+
+def _load_local_vae(dtype: Any) -> Any:
+    path = reqs.local_component("vae", "INLINE_ZIMAGE_VAE")
+    if path is None:
+        return None
+    from diffusers import AutoencoderKL
+
+    if path.is_file():
+        return AutoencoderKL.from_single_file(str(path), torch_dtype=dtype, local_files_only=True)
+    return AutoencoderKL.from_pretrained(str(path), torch_dtype=dtype, local_files_only=True)
+
+
+def _load_local_text_encoder(dtype: Any) -> tuple[Any, Any] | None:
+    path = reqs.local_component("text_encoders", "INLINE_ZIMAGE_TEXT_ENCODER")
+    if path is None or not path.is_dir():  # transformers needs a config dir, not a bare file
+        return None
+    from transformers import AutoModel, AutoTokenizer
+
+    text_encoder = AutoModel.from_pretrained(str(path), torch_dtype=dtype, local_files_only=True)
+    tokenizer = AutoTokenizer.from_pretrained(str(path), local_files_only=True)
+    return text_encoder, tokenizer
+
+
+def _configure(pipe: Any, policy: DevicePolicy) -> None:
+    placement = policy.placement("denoiser")
+    if placement.offload:
+        pipe.enable_model_cpu_offload()  # opt-in only: stream modules on/off the GPU
+    else:
+        pipe.to(str(placement.device))  # default: weights resident on the GPU
+    # Low-VRAM savers that keep weights on the GPU (no offload): slice attention + tile/slice VAE.
+    if policy.attention_slicing():
+        _try(pipe.enable_attention_slicing)
+    if policy.vae_tiling():
+        _try(pipe.enable_vae_tiling)
+        _try(pipe.enable_vae_slicing)
+    _configure_gpu_speed(pipe, placement)
+    _try(pipe.set_progress_bar_config, disable=True)
+
+
+def _configure_gpu_speed(pipe: Any, placement: Placement) -> None:
+    """Throughput tweaks that only apply on a resident-GPU placement (never CPU/offload).
+
+    ``channels_last`` on the conv-based VAE is a safe, default-on win. ``torch.compile`` (the
+    transformer) and xformers attention are opt-in via ``INLINE_COMPILE`` / ``INLINE_XFORMERS`` —
+    both help but have trade-offs (compile warmup, an extra dep), so they stay off by default. The
+    pipeline is cached warm across runs, so a compile cost is paid once. Best-effort."""
+    if placement.offload or placement.device.kind is not DeviceKind.CUDA:
+        return
+    _try(lambda: pipe.vae.to(memory_format=torch.channels_last))
+    if os.environ.get("INLINE_XFORMERS") == "1":
+        _try(pipe.enable_xformers_memory_efficient_attention)
+    if os.environ.get("INLINE_COMPILE") == "1":
+        _try(lambda: setattr(pipe, "transformer", torch.compile(pipe.transformer)))
+
+
+def _torch_dtype(placement: Placement) -> Any:
+    return {
+        DType.FP16: torch.float16,
+        DType.BF16: torch.bfloat16,
+        DType.FP32: torch.float32,
+    }.get(placement.dtype, torch.bfloat16)
+
+
+# --- small helpers ------------------------------------------------------------------------------
+
+
+def _resolve_seed(raw: Any) -> int:
+    """A fixed non-negative seed passes through; -1 (or anything invalid) becomes a fresh random."""
+    try:
+        seed = int(raw)
+    except (TypeError, ValueError):
+        seed = -1
+    return seed if seed >= 0 else random.randint(0, _SEED_MAX)
+
+
+def _load_image(ref: Any) -> Any:
+    from PIL import Image
+
+    if isinstance(ref, AssetRef) and ref.ref == "path" and ref.path:
+        return Image.open(ref.path).convert("RGB")
+    raise ComponentError("Z-Image img2img needs a readable image path input.")
+
+
+def _first(values: list[Any] | None) -> Any:
+    return values[0] if values else None
+
+
+def _first_str(values: list[Any] | None) -> str:
+    value = _first(values)
+    return str(value).strip() if value is not None else ""
+
+
+def _progress(
+    ctx: ExecutionContext,
+    node: Node,
+    phase: Phase,
+    fraction: float,
+    *,
+    step: int | None = None,
+    step_count: int | None = None,
+    status: str = "",
+) -> ProgressEvent:
+    return ProgressEvent(
+        run_id=ctx.run_id,
+        node_id=node.id,
+        phase=phase,
+        fraction=fraction,
+        step=step,
+        step_count=step_count,
+        status=status,
+    )
+
+
+def _try(fn: Any, *args: Any, **kwargs: Any) -> None:
+    """Best-effort optional pipeline tweak; skip if this build lacks it or the op isn't supported
+    here (e.g. xformers not installed, torch.compile unavailable on this device)."""
+    try:
+        fn(*args, **kwargs)
+    except (AttributeError, ValueError, NotImplementedError, RuntimeError, TypeError, ImportError):
+        pass

inline_core/parallel/__init__.py

@@ -0,0 +1,13 @@
+"""Isolated xfuser worker group behind the batched-sampler seam.
+
+A parallel denoise runs in a separate process group the engine spawns and talks to over a local
+socket. The HTTP server, DB, graph, and orchestration stay single-process; only the denoise loop is
+distributed. See PLAN.md (multi-GPU) and `sampling/batch.py` (the seam that routes to the group).
+"""
+
+from __future__ import annotations
+
+from .config import ParallelConfig
+from .group import WorkerGroup, WorkerGroupError
+
+__all__ = ["ParallelConfig", "WorkerGroup", "WorkerGroupError"]

inline_core/parallel/config.py

@@ -0,0 +1,50 @@
+"""What a parallel denoise group is: which model, split how, on which CUDA devices.
+
+The config crosses a process boundary (env var -> worker), so it is JSON-portable. `stub` selects a
+torch-free handler used by the scaffold and the IPC round-trip test.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+
+from ..device.policy import Parallel
+
+# Env vars the manager sets on the worker process (read in worker.main).
+ADDR_ENV = "INLINE_PARALLEL_ADDR"
+CONFIG_ENV = "INLINE_PARALLEL_CONFIG"
+
+
+@dataclass(frozen=True)
+class ParallelConfig:
+    """Identity of a parallel group. One group per (model, parallel, devices) combination."""
+
+    model: str
+    parallel: Parallel = field(default_factory=Parallel)
+    devices: tuple[int, ...] = ()
+    stub: bool = False
+
+    @property
+    def world_size(self) -> int:
+        return self.parallel.world_size
+
+    def to_json(self) -> str:
+        return json.dumps(
+            {
+                "model": self.model,
+                "parallel": asdict(self.parallel),
+                "devices": list(self.devices),
+                "stub": self.stub,
+            }
+        )
+
+    @classmethod
+    def from_json(cls, raw: str) -> ParallelConfig:
+        data = json.loads(raw)
+        return cls(
+            model=data["model"],
+            parallel=Parallel(**data["parallel"]),
+            devices=tuple(data["devices"]),
+            stub=bool(data.get("stub", False)),
+        )

inline_core/parallel/group.py

@@ -0,0 +1,136 @@
+"""The manager side of a parallel denoise group: spawn, handshake, submit, shut down.
+
+One group per parallel model config, spawned lazily on first parallel run and reused. The transport
+is a TCP loopback socket (portable across macOS, Windows, Linux); rank 0 connects back to it. This
+class stays torch-free so the seam is testable with the stub worker and no GPUs.
+"""
+
+from __future__ import annotations
+
+import os
+import socket
+import subprocess
+import sys
+from collections.abc import Callable
+from typing import Any
+
+from ..errors import InlineCoreError
+from .config import ADDR_ENV, CONFIG_ENV, ParallelConfig
+from .launch import Launcher, default_launcher
+from .protocol import MessageType, recv_message, send_message
+
+ProgressHandler = Callable[[int, int], None]
+
+
+class WorkerGroupError(InlineCoreError):
+    """The worker group failed to start, died, or returned an error for a job."""
+
+
+class WorkerGroup:
+    def __init__(
+        self,
+        config: ParallelConfig,
+        launcher: Launcher | None = None,
+        *,
+        ready_timeout: float = 120.0,
+        shutdown_timeout: float = 10.0,
+    ) -> None:
+        self._config = config
+        self._launcher = launcher or default_launcher(config)
+        self._ready_timeout = ready_timeout
+        self._shutdown_timeout = shutdown_timeout
+        self._listener: socket.socket | None = None
+        self._conn: socket.socket | None = None
+        self._process: subprocess.Popen[bytes] | None = None
+
+    def start(self) -> None:
+        listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        listener.bind(("127.0.0.1", 0))
+        listener.listen(1)
+        host, port = listener.getsockname()
+        self._listener = listener
+        self._process = subprocess.Popen(
+            self._launcher.command(self._config), env=self._child_env(host, port)
+        )
+        listener.settimeout(self._ready_timeout)
+        try:
+            conn, _ = listener.accept()
+        except OSError as exc:
+            self._terminate()
+            raise WorkerGroupError("worker group did not connect in time") from exc
+        conn.settimeout(None)
+        self._conn = conn
+        ready = recv_message(conn)
+        if ready is None or ready.get("type") != MessageType.READY:
+            self._terminate()
+            raise WorkerGroupError("worker group did not report ready")
+
+    def submit(
+        self, payload: dict[str, Any], on_progress: ProgressHandler | None = None
+    ) -> dict[str, Any]:
+        conn = self._require_conn()
+        send_message(conn, {"type": MessageType.JOB, "payload": payload})
+        while True:
+            message = recv_message(conn)
+            if message is None:
+                raise WorkerGroupError("worker group closed the connection mid-job")
+            kind = message.get("type")
+            if kind == MessageType.PROGRESS:
+                if on_progress is not None:
+                    on_progress(int(message["step"]), int(message["total"]))
+            elif kind == MessageType.RESULT:
+                result: dict[str, Any] = message["payload"]
+                return result
+            elif kind == MessageType.ERROR:
+                raise WorkerGroupError(str(message.get("message", "worker group error")))
+            else:
+                raise WorkerGroupError(f"unexpected message type {kind!r}")
+
+    def shutdown(self) -> None:
+        if self._conn is not None:
+            try:
+                send_message(self._conn, {"type": MessageType.SHUTDOWN})
+            except OSError:
+                pass
+        self._terminate()
+
+    def _require_conn(self) -> socket.socket:
+        if self._conn is None:
+            raise WorkerGroupError("worker group is not started")
+        return self._conn
+
+    def _child_env(self, host: str, port: int) -> dict[str, str]:
+        env = dict(os.environ)
+        env[ADDR_ENV] = f"{host}:{port}"
+        env[CONFIG_ENV] = self._config.to_json()
+        # Let the worker import inline_core under an editable/src layout, where the package is on
+        # sys.path but not on the inherited PYTHONPATH. Harmless once pip-installed.
+        paths = [p for p in sys.path if p]
+        existing = env.get("PYTHONPATH")
+        if existing:
+            paths = [existing, *paths]
+        env["PYTHONPATH"] = os.pathsep.join(paths)
+        return env
+
+    def _terminate(self) -> None:
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None
+        if self._listener is not None:
+            self._listener.close()
+            self._listener = None
+        process = self._process
+        if process is not None:
+            try:
+                process.wait(timeout=self._shutdown_timeout)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                process.wait()
+            self._process = None
+
+    def __enter__(self) -> WorkerGroup:
+        self.start()
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.shutdown()

inline_core/parallel/launch.py

@@ -0,0 +1,44 @@
+"""How the worker process group is launched. The transport and protocol do not care which.
+
+DirectLauncher runs one process (rank 0) for a single GPU, CPU, or the stub test. TorchrunLauncher
+starts N ranks via torchrun for a real multi-GPU split; its exact flags are finalized against the
+2-GPU box when the xfuser handler lands (C2). The manager selects by world size.
+"""
+
+from __future__ import annotations
+
+import sys
+from abc import ABC, abstractmethod
+
+from .config import ParallelConfig
+
+_WORKER_MODULE = "inline_core.parallel.worker"
+
+
+class Launcher(ABC):
+    @abstractmethod
+    def command(self, config: ParallelConfig) -> list[str]:
+        """The argv that starts the worker process group."""
+
+
+class DirectLauncher(Launcher):
+    def command(self, config: ParallelConfig) -> list[str]:
+        return [sys.executable, "-m", _WORKER_MODULE]
+
+
+class TorchrunLauncher(Launcher):
+    def command(self, config: ParallelConfig) -> list[str]:
+        return [
+            sys.executable,
+            "-m",
+            "torch.distributed.run",
+            f"--nproc_per_node={config.world_size}",
+            "--rdzv-backend=c10d",
+            "--rdzv-endpoint=127.0.0.1:0",
+            "-m",
+            _WORKER_MODULE,
+        ]
+
+
+def default_launcher(config: ParallelConfig) -> Launcher:
+    return TorchrunLauncher() if config.world_size > 1 else DirectLauncher()

inline_core/parallel/protocol.py

@@ -0,0 +1,53 @@
+"""Framed JSON messages over a stream socket. Length-prefixed so reads never split a message.
+
+Control only: jobs, progress, results, shutdown. Large latent/image tensors move over a side channel
+(shared memory) added with the xfuser handler; the control plane stays small and JSON-portable.
+"""
+
+from __future__ import annotations
+
+import json
+import socket
+import struct
+from typing import Any
+
+_HEADER = struct.Struct(">I")  # 4-byte big-endian body length
+
+
+class MessageType:
+    READY = "ready"  # worker -> manager: connected and serving
+    JOB = "job"  # manager -> worker: run this payload
+    PROGRESS = "progress"  # worker -> manager: step of total
+    RESULT = "result"  # worker -> manager: job payload result
+    ERROR = "error"  # worker -> manager: job failed, message attached
+    SHUTDOWN = "shutdown"  # manager -> worker: stop serving and exit
+
+
+def send_message(sock: socket.socket, message: dict[str, Any]) -> None:
+    body = json.dumps(message).encode("utf-8")
+    sock.sendall(_HEADER.pack(len(body)) + body)
+
+
+def recv_message(sock: socket.socket) -> dict[str, Any] | None:
+    """The next message, or None when the peer closed the connection cleanly."""
+    header = _recv_exactly(sock, _HEADER.size)
+    if header is None:
+        return None
+    (length,) = _HEADER.unpack(header)
+    body = _recv_exactly(sock, length)
+    if body is None:
+        return None
+    decoded: dict[str, Any] = json.loads(body.decode("utf-8"))
+    return decoded
+
+
+def _recv_exactly(sock: socket.socket, count: int) -> bytes | None:
+    chunks: list[bytes] = []
+    remaining = count
+    while remaining > 0:
+        chunk = sock.recv(remaining)
+        if not chunk:
+            return None
+        chunks.append(chunk)
+        remaining -= len(chunk)
+    return b"".join(chunks)