inline-core 1.1.1__py3-none-any.whl

inline_core/parallel/registry.py

@@ -0,0 +1,31 @@
+"""One running worker group per parallel config, created lazily and reused across jobs.
+
+The server owns a single registry and shuts it down with the app; the parallel sampler asks it for
+the group that matches a job's placement. Keying on the frozen ParallelConfig means the same model
+and split reuse one process group.
+"""
+
+from __future__ import annotations
+
+from .config import ParallelConfig
+from .group import WorkerGroup
+from .launch import Launcher
+
+
+class GroupRegistry:
+    def __init__(self, launcher: Launcher | None = None) -> None:
+        self._launcher = launcher
+        self._groups: dict[ParallelConfig, WorkerGroup] = {}
+
+    def get_or_create(self, config: ParallelConfig) -> WorkerGroup:
+        group = self._groups.get(config)
+        if group is None:
+            group = WorkerGroup(config, self._launcher)
+            group.start()
+            self._groups[config] = group
+        return group
+
+    def shutdown_all(self) -> None:
+        while self._groups:
+            _, group = self._groups.popitem()
+            group.shutdown()

inline_core/parallel/worker.py

@@ -0,0 +1,75 @@
+"""Rank entrypoint for a parallel denoise group. Launched by a Launcher, one process per rank.
+
+Rank 0 owns the IPC channel to the manager: it receives jobs, streams progress, and returns results
+through a pluggable handler. The handler decides what a job means: the stub handler (no torch) backs
+the scaffold and the round-trip test; the xfuser handler (loads the sharded pipeline and runs the
+collective denoise across all ranks) lands with the Z-Image runner (C2).
+"""
+
+from __future__ import annotations
+
+import os
+import socket
+from collections.abc import Callable
+from typing import Any
+
+from .config import ADDR_ENV, CONFIG_ENV, ParallelConfig
+from .protocol import MessageType, recv_message, send_message
+
+# report(step, total) streams progress to the manager while a job runs.
+ProgressFn = Callable[[int, int], None]
+# handler(config, payload, report) -> result payload.
+Handler = Callable[[ParallelConfig, dict[str, Any], ProgressFn], dict[str, Any]]
+
+
+def _stub_handler(
+    config: ParallelConfig, payload: dict[str, Any], report: ProgressFn
+) -> dict[str, Any]:
+    total = int(payload.get("steps", 1))
+    for step in range(1, total + 1):
+        report(step, total)
+    return {"echo": payload, "model": config.model, "world_size": config.world_size}
+
+
+def _select_handler(config: ParallelConfig) -> Handler:
+    if config.stub:
+        return _stub_handler
+    raise NotImplementedError(
+        "the xfuser denoise handler lands with the Z-Image runner (C2); "
+        "only the stub handler is available today"
+    )
+
+
+def serve(sock: socket.socket, config: ParallelConfig, handler: Handler) -> None:
+    send_message(sock, {"type": MessageType.READY, "worldSize": config.world_size})
+    while True:
+        message = recv_message(sock)
+        if message is None or message.get("type") == MessageType.SHUTDOWN:
+            return
+        if message.get("type") != MessageType.JOB:
+            continue
+        payload: dict[str, Any] = message.get("payload", {})
+
+        def report(step: int, total: int) -> None:
+            send_message(sock, {"type": MessageType.PROGRESS, "step": step, "total": total})
+
+        try:
+            result = handler(config, payload, report)
+        except Exception as exc:  # translate into a protocol error; the manager re-raises typed
+            send_message(sock, {"type": MessageType.ERROR, "message": str(exc)})
+            continue
+        send_message(sock, {"type": MessageType.RESULT, "payload": result})
+
+
+def main() -> None:
+    config = ParallelConfig.from_json(os.environ[CONFIG_ENV])
+    host, _, port = os.environ[ADDR_ENV].partition(":")
+    sock = socket.create_connection((host, int(port)))
+    try:
+        serve(sock, config, _select_handler(config))
+    finally:
+        sock.close()
+
+
+if __name__ == "__main__":
+    main()

inline_core/runtime/__init__.py

@@ -0,0 +1 @@
+"""Per-run execution: the progress event bus, run state, and the execution context."""

inline_core/runtime/context.py

@@ -0,0 +1,31 @@
+"""The per-run execution context threaded through every component call."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from threading import Event
+
+from ..device.policy import DevicePolicy
+from .progress import ProgressEmitter
+
+
+class CancelToken:
+    """Cooperative cancellation. The executor checks it between nodes and steps."""
+
+    def __init__(self) -> None:
+        self._event = Event()
+
+    def cancel(self) -> None:
+        self._event.set()
+
+    @property
+    def cancelled(self) -> bool:
+        return self._event.is_set()
+
+
+@dataclass
+class ExecutionContext:
+    run_id: str
+    policy: DevicePolicy
+    emitter: ProgressEmitter
+    cancel: CancelToken

inline_core/runtime/file_store.py

@@ -0,0 +1,51 @@
+"""A file-backed take store: writes a decoded image to <root>/<take_id>.png and records the take."""
+
+from __future__ import annotations
+
+import hashlib
+import time
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from ..media import MediaKind
+from ..takes import Take
+from .store import TakeStore
+
+
+class FileTakeStore(TakeStore):
+    def __init__(self, root: Path) -> None:
+        self._root = root
+
+    def save(self, run_id: str, node_id: str, image: Any, params: dict[str, Any]) -> Take:
+        self._root.mkdir(parents=True, exist_ok=True)
+        take_id = f"take_{uuid4().hex[:12]}"
+        path = self._root / f"{take_id}.png"
+        _to_pil(image).save(path, format="PNG")
+        data = path.read_bytes()
+        return Take(
+            id=take_id,
+            run_id=run_id,
+            node_id=node_id,
+            kind=MediaKind.IMAGE,
+            uri=str(path),
+            hash=f"sha256-{hashlib.sha256(data).hexdigest()}",
+            params=dict(params),
+            created_at=int(time.time() * 1000),
+        )
+
+
+def _to_pil(image: Any) -> Any:
+    from PIL import Image
+
+    if isinstance(image, Image.Image):
+        return image
+    if hasattr(image, "detach"):  # a torch tensor
+        image = image.detach().to("cpu").numpy()
+    import numpy as np
+
+    array = np.asarray(image)
+    if array.dtype != np.uint8:
+        scaled = array * 255.0 if float(array.max(initial=0.0)) <= 1.0 else array
+        array = scaled.clip(0, 255).round().astype(np.uint8)
+    return Image.fromarray(array)

inline_core/runtime/progress.py

@@ -0,0 +1,83 @@
+"""Run events and the emitter seam. Mirrors the websocket events in docs/contract.md section 6.
+
+Coalescing (the contract's bounded event rate) is a wrapper added at the serving layer; the executor
+and components just emit.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+
+from ..takes import Take
+
+
+class Phase(str, Enum):
+    QUEUED = "queued"
+    PREPARING = "preparing"
+    LOADING = "loading"
+    ENCODE = "encode"
+    SAMPLE = "sample"
+    DECODE = "decode"
+    SAVE = "save"
+
+
+@dataclass(frozen=True)
+class ProgressEvent:
+    run_id: str
+    node_id: str
+    phase: Phase
+    fraction: float
+    step: int | None = None
+    step_count: int | None = None
+    eta_ms: int | None = None
+    status: str = ""
+
+
+@dataclass(frozen=True)
+class NodeDoneEvent:
+    run_id: str
+    node_id: str
+    cached: bool
+    takes: list[Take] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class RunDoneEvent:
+    run_id: str
+
+
+@dataclass(frozen=True)
+class CancelledEvent:
+    run_id: str
+
+
+@dataclass(frozen=True)
+class ErrorEvent:
+    run_id: str
+    message: str
+    node_id: str | None = None
+
+
+RunEvent = ProgressEvent | NodeDoneEvent | RunDoneEvent | CancelledEvent | ErrorEvent
+
+
+class ProgressEmitter(ABC):
+    @abstractmethod
+    def emit(self, event: RunEvent) -> None: ...
+
+
+class NullEmitter(ProgressEmitter):
+    def emit(self, event: RunEvent) -> None:
+        return None
+
+
+class CollectingEmitter(ProgressEmitter):
+    """Keeps every event in order. For tests and the snapshot builder."""
+
+    def __init__(self) -> None:
+        self.events: list[RunEvent] = []
+
+    def emit(self, event: RunEvent) -> None:
+        self.events.append(event)

inline_core/runtime/run.py

@@ -0,0 +1,113 @@
+"""Run state: the durable snapshot GET /v1/runs/{id} serves, kept fresh by applying run events."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+from ..takes import Take
+from .progress import (
+    CancelledEvent,
+    ErrorEvent,
+    NodeDoneEvent,
+    ProgressEmitter,
+    ProgressEvent,
+    RunDoneEvent,
+    RunEvent,
+)
+
+
+class RunStatus(str, Enum):
+    QUEUED = "queued"
+    RUNNING = "running"
+    DONE = "done"
+    ERROR = "error"
+    CANCELLED = "cancelled"
+
+
+class NodeState(str, Enum):
+    QUEUED = "queued"
+    RUNNING = "running"
+    CACHED = "cached"
+    DONE = "done"
+    ERROR = "error"
+
+
+@dataclass
+class NodeRuntimeState:
+    state: NodeState = NodeState.QUEUED
+    phase: str | None = None
+    fraction: float = 0.0
+    step: int | None = None
+    step_count: int | None = None
+    status: str = ""
+
+
+@dataclass
+class RunError:
+    message: str
+    node_id: str | None = None
+
+
+@dataclass
+class RunState:
+    run_id: str
+    target: str
+    status: RunStatus = RunStatus.QUEUED
+    fraction: float = 0.0
+    nodes: dict[str, NodeRuntimeState] = field(default_factory=dict)
+    takes: list[Take] = field(default_factory=list)
+    error: RunError | None = None
+
+
+_COMPLETE = (NodeState.DONE, NodeState.CACHED)
+
+
+def _recompute_fraction(state: RunState) -> None:
+    if not state.nodes:
+        state.fraction = 0.0
+        return
+    total = sum(1.0 if n.state in _COMPLETE else n.fraction for n in state.nodes.values())
+    state.fraction = total / len(state.nodes)
+
+
+def apply_event(state: RunState, event: RunEvent) -> None:
+    """Fold a run event into the run state. The single event -> snapshot mapping."""
+    if isinstance(event, ProgressEvent):
+        node = state.nodes.setdefault(event.node_id, NodeRuntimeState())
+        node.state = NodeState.RUNNING
+        node.phase = event.phase.value
+        node.fraction = event.fraction
+        node.step = event.step
+        node.step_count = event.step_count
+        node.status = event.status
+        _recompute_fraction(state)
+    elif isinstance(event, NodeDoneEvent):
+        node = state.nodes.setdefault(event.node_id, NodeRuntimeState())
+        node.state = NodeState.CACHED if event.cached else NodeState.DONE
+        node.fraction = 1.0
+        known = {t.id for t in state.takes}
+        state.takes.extend(t for t in event.takes if t.id not in known)
+        _recompute_fraction(state)
+    elif isinstance(event, RunDoneEvent):
+        state.status = RunStatus.DONE
+        state.fraction = 1.0
+    elif isinstance(event, CancelledEvent):
+        state.status = RunStatus.CANCELLED
+    elif isinstance(event, ErrorEvent):
+        state.status = RunStatus.ERROR
+        state.error = RunError(message=event.message, node_id=event.node_id)
+        if event.node_id is not None:
+            state.nodes.setdefault(event.node_id, NodeRuntimeState()).state = NodeState.ERROR
+
+
+class StateTrackingEmitter(ProgressEmitter):
+    """Applies each event to a RunState, then forwards it to the real emitter."""
+
+    def __init__(self, delegate: ProgressEmitter, state: RunState) -> None:
+        self._delegate = delegate
+        self._state = state
+
+    def emit(self, event: RunEvent) -> None:
+        apply_event(self._state, event)
+        self._delegate.emit(event)

inline_core/runtime/store.py

@@ -0,0 +1,18 @@
+"""The take store seam: persist a decoded output as an immutable take (bytes, hash, uri).
+
+Phase 1's implementation writes into the project's takes/ folder. Kept behind this interface so a
+fleet object store swaps in without touching the executor.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from ..takes import Take
+
+
+class TakeStore(ABC):
+    @abstractmethod
+    def save(self, run_id: str, node_id: str, image: Any, params: dict[str, Any]) -> Take:
+        """Persist a decoded image (PIL, numpy, or tensor) as an immutable take."""

inline_core/sampling/__init__.py

@@ -0,0 +1 @@
+"""The batched-sampler boundary: the graph submits jobs, this batches and steps them on the GPU."""

inline_core/sampling/batch.py

@@ -0,0 +1,116 @@
+"""The graph/GPU boundary. The graph never runs the denoise loop inline; it submits a SampleJob.
+
+Phase 1 runs one job at a time behind this seam. Phase 5 adds cross-request batching (grouping by
+model family, resolution, and adapter bucket) without changing the interface. Multi-GPU splits one
+job across a worker group (XFuserBatchedSampler) when the device policy asks for a parallel denoise.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Protocol
+
+from ..components.conditioning import Conditioning, Latents
+from ..components.interfaces import Denoiser, Sampler, Scheduler, StepCallback, StepInfo
+
+if TYPE_CHECKING:
+    from ..device.policy import Placement
+    from ..parallel.config import ParallelConfig
+    from ..parallel.group import ProgressHandler
+    from ..parallel.registry import GroupRegistry
+    from ..runtime.context import ExecutionContext
+
+
+@dataclass
+class SampleJob:
+    """One request to sample. Compatible jobs are grouped and stepped together."""
+
+    denoiser: Denoiser
+    scheduler: Scheduler
+    sampler: Sampler
+    latents: Latents
+    conditioning: Conditioning
+    steps: int
+    on_step: StepCallback | None = None
+
+
+class BatchedSampler(ABC):
+    @abstractmethod
+    def submit(self, job: SampleJob, ctx: ExecutionContext) -> Latents: ...
+
+
+class InlineBatchedSampler(BatchedSampler):
+    """Phase 1: no batching. Run the job through its sampler directly."""
+
+    def submit(self, job: SampleJob, ctx: ExecutionContext) -> Latents:
+        return job.sampler.sample(
+            job.denoiser,
+            job.scheduler,
+            job.latents,
+            job.conditioning,
+            steps=job.steps,
+            ctx=ctx,
+            on_step=job.on_step,
+        )
+
+
+class ParallelCodec(Protocol):
+    """Bridges a SampleJob and its worker group: which group, request out, latents back.
+
+    The real codec moves torch tensors over a side channel and builds the xfuser request; it lives
+    with the model runner (C2) so this module stays torch-free and the group stays mockable.
+    """
+
+    def config(self, placement: Placement) -> ParallelConfig:
+        """The group identity (model + split) for this placement."""
+        ...
+
+    def to_request(self, job: SampleJob) -> dict[str, Any]:
+        """The serializable request the worker runs against its resident pipeline."""
+        ...
+
+    def from_result(self, result: dict[str, Any], job: SampleJob) -> Latents:
+        """The sampled latents the worker returned."""
+        ...
+
+
+class XFuserBatchedSampler(BatchedSampler):
+    """Routes a parallel-placement denoise to the worker group; everything else runs inline.
+
+    The device policy decides: a single-GPU or CPU run has no parallel placement, so it keeps the
+    in-process path with zero overhead. Only a multi-GPU denoiser placement crosses to the group.
+    """
+
+    def __init__(
+        self,
+        registry: GroupRegistry,
+        codec: ParallelCodec,
+        inline: BatchedSampler | None = None,
+    ) -> None:
+        self._registry = registry
+        self._codec = codec
+        self._inline = inline or InlineBatchedSampler()
+
+    def submit(self, job: SampleJob, ctx: ExecutionContext) -> Latents:
+        placement = ctx.policy.placement("denoiser")
+        if placement.parallel is None:
+            return self._inline.submit(job, ctx)
+        group = self._registry.get_or_create(self._codec.config(placement))
+        result = group.submit(self._codec.to_request(job), _forward_progress(job))
+        return self._codec.from_result(result, job)
+
+    def close(self) -> None:
+        self._registry.shutdown_all()
+
+
+def _forward_progress(job: SampleJob) -> ProgressHandler | None:
+    """Feed worker step ticks into the job's on_step callback, so streaming stays unchanged."""
+    on_step = job.on_step
+    if on_step is None:
+        return None
+
+    def report(step: int, total: int) -> None:
+        on_step(StepInfo(step=step, total=total))
+
+    return report

inline_core/server/__init__.py

@@ -0,0 +1 @@
+"""The /v1 HTTP and websocket serving layer. A thin shell over the graph engine (see contract)."""

inline_core/server/__main__.py

@@ -0,0 +1,61 @@
+"""Run the engine: `python -m inline_core.server`. Registers models whose deps are installed."""
+
+from __future__ import annotations
+
+import uvicorn
+
+from ..config import data_dir, server_host, server_port
+from ..device.memory import MemoryPolicy
+from ..graph.cache import InMemoryCache
+from ..graph.registry import build_default_registry
+from ..runtime.file_store import FileTakeStore
+from ..studio import config as studio_config
+from ..studio.store import StudioStore
+from .app import create_app
+from .bootstrap import register_models
+from .frontend import resolve_frontend_root
+from .rpc import EventBroadcaster, RpcRouter
+from .run_store import SqliteRunStore
+
+
+def main() -> None:
+    policy = MemoryPolicy()
+    registry = build_default_registry()
+    data = data_dir()
+    takes = data / "takes"
+    store = FileTakeStore(takes)
+    run_store = SqliteRunStore(data / "runs.db")
+    registered = register_models(registry, store, policy)
+    print(f"Registered models: {registered or 'none (source nodes only)'}")
+    frontend_root = resolve_frontend_root()
+    fe = frontend_root or "none (API only); use --front-end-root or install the frontend package"
+    print(f"Frontend: {fe}")
+    # The Studio app-backend: Core is the sole native backend (projects, frames, moodboard, assets,
+    # generation, fal, timeline). Every InlineStudioApi channel is handled here.
+    rpc = RpcRouter()
+    events = EventBroadcaster()
+    store = StudioStore(
+        studio_config.data_dir(),
+        studio_config.workspace_dir(),
+        default_comfy_url=studio_config.DEFAULT_COMFY_URL,
+        default_core_url=studio_config.DEFAULT_CORE_URL,
+    )
+    print(f"Studio data: {studio_config.data_dir()}  |  workspace: {studio_config.workspace_dir()}")
+    app = create_app(
+        registry=registry,
+        cache=InMemoryCache(),
+        policy=policy,
+        run_store=run_store,
+        takes_dir=str(takes),
+        frontend_root=frontend_root,
+        rpc=rpc,
+        events=events,
+        studio_store=store,
+    )
+    host, port = server_host(), server_port()
+    print(f"Serving on http://{host}:{port}")
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    main()