cluster: parts of a video pipeline on other machines

Per-model 'components' (text_encoder / low_noise / vae → cluster node): the
prompt is encoded where the text encoder is, the Wan 2.2 low-noise expert
takes over at the expert boundary with the latents handed over, the VAE
decodes where it lives and returns the finished video. The far side runs
the ordinary generation endpoint with a _handoff block, so loading,
acceleration and LoRAs happen as usual there; the scheduler is sliced to
resume mid-loop (codai/cluster/components.py). Models page block + tests.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01FKm7Z56iKjHgNkHc6VhWKz

codai/admin/routes.py

@@ -2778,6 +2778,16 @@ async def api_model_configure(request: Request, username: str = Depends(require_
             entry["distribute"] = blk
         else:
             entry.pop("distribute", None)
+    # Parts of a video pipeline on other machines (codai/cluster/components.py).
+    if "components" in data:
+        src = data.get("components") if isinstance(data.get("components"), dict) else {}
+        blk = {k: str(v).strip() for k, v in src.items()
+               if k in ("text_encoder", "low_noise", "vae") and isinstance(v, str)
+               and v.strip() and v.strip().lower() not in ("local", "here", "auto")}
+        if blk:
+            entry["components"] = blk
+        else:
+            entry.pop("components", None)
     # Per-model engine blocks laid over the global template: vLLM (parallelism,
     # ray nodes), SGLang/kt (nnodes, ranks' commands). A block with nothing set
     # is dropped so the entry stays clean.

codai/admin/templates/models.html

@@ -637,6 +637,15 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
           </div>
         </div>
       </div>
+      <div class="form-row" id="cfg-comp-row" style="margin-top:.75rem;display:none;border:1px solid var(--border,#333);border-radius:8px;padding:.6rem">
+        <label class="form-label" style="font-weight:600">Pipeline parts on other machines <span class="muted" style="font-weight:400">— video models (Wan)</span></label>
+        <span class="form-hint" style="font-size:11px;line-height:1.5">A video model is several models run one after the other; what passes between them is small. Name a cluster node for a part and that part runs there: the <b>text encoder</b> (UMT5-XXL, 11 GB) encodes the prompt elsewhere; the <b>low-noise expert</b> of a Wan 2.2 dual-expert model takes over at the expert boundary (latents handed over, a few MB); the <b>VAE</b> decodes elsewhere and returns the finished video. Blank = here. The node needs the same model.</span>
+        <div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:.6rem;margin-top:.4rem">
+          <div class="form-row" style="margin:0"><label class="form-label">Text encoder on</label><input type="text" id="cfg-comp-te" class="form-input" placeholder="here"></div>
+          <div class="form-row" style="margin:0"><label class="form-label">Low-noise expert on</label><input type="text" id="cfg-comp-low" class="form-input" placeholder="here"></div>
+          <div class="form-row" style="margin:0"><label class="form-label">VAE decode on</label><input type="text" id="cfg-comp-vae" class="form-input" placeholder="here"></div>
+        </div>
+      </div>
       <div class="form-row" id="cfg-rpc-row" style="margin-top:.75rem;display:none;border:1px solid var(--border,#333);border-radius:8px;padding:.6rem">
         <label class="form-label" style="font-weight:600">Cards on other machines <span class="muted" style="font-weight:400">— llama.cpp RPC servers, one model over several machines</span></label>
         <input type="text" id="cfg-rpc-servers" class="form-input" list="cfg-rpc-known" placeholder="host:50052, host2:50052   (blank = this machine only)">
@@ -4344,6 +4353,10 @@ function _updateDistVisibility(){
 function _populateCluster(m, s){
   window._cfgCurrentPath = (m && (m.path || m.id)) || '';
   const set = (id,v)=>{ const el=document.getElementById(id); if(el!=null) el.value=(v==null?'':v); };
+  const cp = s.components || {};
+  set('cfg-comp-te', cp.text_encoder || ''); set('cfg-comp-low', cp.low_noise || ''); set('cfg-comp-vae', cp.vae || '');
+  const crow = document.getElementById('cfg-comp-row');
+  if (crow) { const isVideo = [...document.querySelectorAll('.cfg-type-cb')].some(cb => cb.checked && cb.value === 'video_models') || !!cp.text_encoder || !!cp.low_noise || !!cp.vae; crow.style.display = isVideo ? '' : 'none'; }
   const dz = s.distribute || {};
   document.getElementById('cfg-dist-enabled').checked = !!dz.enabled;
   set('cfg-dist-nodes', Array.isArray(dz.nodes) ? dz.nodes.join(', ') : (dz.nodes || ''));
@@ -4376,6 +4389,7 @@ function _collectCluster(){
   const out = {
     node_path: val('cfg-node-path'),
     rpc_servers: val('cfg-rpc-servers'),
+    components: { text_encoder: val('cfg-comp-te'), low_noise: val('cfg-comp-low'), vae: val('cfg-comp-vae') },
     distribute: {
       enabled: !!document.getElementById('cfg-dist-enabled').checked,
       nodes: val('cfg-dist-nodes'), max_parts: num('cfg-dist-max'),

codai/api/video.py

@@ -37,6 +37,7 @@ import uuid
 from typing import List, Optional
 
 from fastapi import APIRouter, HTTPException, Request
+from fastapi.responses import JSONResponse
 
 from codai.models.manager import multi_model_manager
 from codai.pydantic.videorequest import (
@@ -2161,6 +2162,161 @@ def _build_vace_conditioning(cond_frames, num_frames: int, width: int, height: i
     return video, mask
 
 
+class _HandoffPayload:
+    """What a component run answers instead of frames: the endpoint returns
+    it as JSON to the machine that asked (latents, embeddings, or the video
+    when this machine held the VAE)."""
+
+    def __init__(self, payload: dict):
+        self.payload = payload
+
+
+def _wan_embeds_into_kw(pipe, kw: dict, pe_b64: str, ne_b64: str) -> None:
+    from codai.cluster.components import b64_to_tensor
+    import torch as _torch
+    dev = pipe._execution_device
+    dt = getattr(pipe, 'transformer', None)
+    dt = getattr(dt, 'dtype', None) or _torch.float32
+    kw['prompt_embeds'] = b64_to_tensor(pe_b64, device=dev, dtype=dt)
+    if ne_b64:
+        kw['negative_prompt_embeds'] = b64_to_tensor(ne_b64, device=dev, dtype=dt)
+    kw.pop('prompt', None)
+    kw.pop('negative_prompt', None)
+
+
+def _run_handoff(pipe, kw: dict, request, h: dict):
+    """This machine's part of a run that started elsewhere (see components.py)."""
+    from codai.cluster.components import SlicedTimesteps, b64_to_tensor, tensor_to_b64
+    import torch as _torch
+    op = str(h.get('op'))
+    if op == 'encode':
+        dev = pipe._execution_device
+        pe, ne = pipe.encode_prompt(
+            prompt=kw.get('prompt') or '', negative_prompt=kw.get('negative_prompt'),
+            do_classifier_free_guidance=True, num_videos_per_prompt=1,
+            max_sequence_length=int(kw.get('max_sequence_length') or 512), device=dev)
+        print(f"  [video][components] encoded the prompt for {request.model}", flush=True)
+        return _HandoffPayload({'prompt_embeds': tensor_to_b64(pe),
+                                'negative_prompt_embeds': tensor_to_b64(ne) if ne is not None else ''})
+    if h.get('prompt_embeds'):
+        _wan_embeds_into_kw(pipe, kw, h['prompt_embeds'], h.get('negative_prompt_embeds') or '')
+    dt = getattr(getattr(pipe, 'transformer', None), 'dtype', None) or _torch.float32
+    latents = b64_to_tensor(h['latents'], device=pipe._execution_device, dtype=dt)
+    kw['latents'] = latents
+    steps = int(kw.get('num_inference_steps') or 25)
+    if op == 'decode':
+        start = steps                              # no denoising: straight to the VAE
+    elif op == 'denoise':
+        start = int(h.get('start_step') or 0)
+    else:
+        raise ValueError(f"unknown handoff op {op!r}")
+    want_latent = (h.get('return') == 'latent')
+    if want_latent:
+        kw['output_type'] = 'latent'
+    print(f"  [video][components] {op} from step {start}/{steps} for {request.model}"
+          f"{' → latents' if want_latent else ' → video'}", flush=True)
+    with SlicedTimesteps(pipe.scheduler, start):
+        result = pipe(**kw)
+    if want_latent:
+        out = getattr(result, 'frames', None)
+        if out is None:
+            out = result[0]
+        return _HandoffPayload({'latents': tensor_to_b64(out), 'start_step': start})
+    return _frames_from_result(result)
+
+
+def _run_pipeline_split(pipe, kw: dict, request, comp: dict):
+    """The head's side: encode where the text encoder is, denoise up to the
+    expert boundary here, hand the latents to the low-noise machine, decode
+    where the VAE is. Returns frames, or the far machine's finished answer."""
+    from codai.cluster.components import (SlicedTimesteps, b64_to_tensor, boundary_step,
+                                          call_handoff, request_json, tensor_to_b64)
+    import torch as _torch
+    req_json = request_json(request)
+    steps = int(kw.get('num_inference_steps') or 25)
+    embeds = {}
+    if comp.get('text_encoder'):
+        r = call_handoff(comp['text_encoder'], req_json, {'op': 'encode'})
+        embeds = {'prompt_embeds': r['prompt_embeds'],
+                  'negative_prompt_embeds': r.get('negative_prompt_embeds') or ''}
+        _wan_embeds_into_kw(pipe, kw, embeds['prompt_embeds'], embeds['negative_prompt_embeds'])
+        print(f"  [video][components] prompt encoded on {comp['text_encoder']}", flush=True)
+    k = boundary_step(pipe, kw) if comp.get('low_noise') else None
+    if k is not None and k <= 0:
+        k = None                     # every step is low-noise: nothing to do here
+    vae_remote = comp.get('vae')
+    if k is None and not vae_remote:
+        return _run_pipeline(pipe, kw)
+
+    # Run here up to the boundary (or the whole loop), keeping latents.
+    user_cb = kw.get('callback_on_step_end')
+    captured = {}
+
+    def _cb(p, i, t, cb_kwargs):
+        out = user_cb(p, i, t, cb_kwargs) if user_cb else cb_kwargs
+        if k is not None and i + 1 >= k:
+            captured['latents'] = (out or cb_kwargs).get('latents', cb_kwargs.get('latents'))
+            p._interrupt = True
+        return out
+    kw['callback_on_step_end'] = _cb
+    kw['callback_on_step_end_tensor_inputs'] = ['latents']
+    kw['output_type'] = 'latent'
+    pipe._interrupt = False
+    try:
+        result = pipe(**kw)
+    finally:
+        pipe._interrupt = False
+    latents = captured.get('latents')
+    if latents is None:
+        latents = getattr(result, 'frames', None)
+        if latents is None:
+            latents = result[0]
+    kw['callback_on_step_end'] = user_cb
+    kw.pop('callback_on_step_end_tensor_inputs', None)
+    kw.pop('output_type', None)
+
+    if k is not None:
+        # The low-noise expert lives elsewhere: hand over and, when the VAE is
+        # there too, let it finish the video.
+        handoff = {'op': 'denoise', 'start_step': k, 'latents': tensor_to_b64(latents),
+                   'return': 'latent' if not vae_remote else 'video', **embeds}
+        print(f"  [video][components] steps {k}..{steps} on {comp['low_noise']}"
+              f"{' with its VAE' if vae_remote else ''}", flush=True)
+        r = call_handoff(comp['low_noise'], req_json, handoff)
+        if vae_remote:
+            return _HandoffPayload(r)
+        latents = b64_to_tensor(r['latents'], device=pipe._execution_device,
+                                dtype=getattr(pipe.transformer, 'dtype', _torch.float32))
+    elif vae_remote:
+        print(f"  [video][components] decoding on {vae_remote}", flush=True)
+        r = call_handoff(vae_remote, req_json, {'op': 'decode', 'latents': tensor_to_b64(latents),
+                                                'return': 'video', **embeds})
+        return _HandoffPayload(r)
+
+    # Decode here: a zero-step run straight into the VAE and the video processor.
+    kw['latents'] = latents
+    with SlicedTimesteps(pipe.scheduler, steps):
+        result = pipe(**kw)
+    return _frames_from_result(result)
+
+
+def _frames_from_result(result):
+    frames_raw = getattr(result, 'frames', None)
+    if frames_raw is None:
+        frames_raw = result[0]
+    if isinstance(frames_raw, list):
+        if frames_raw and isinstance(frames_raw[0], list):
+            return frames_raw[0]
+        return frames_raw
+    try:
+        import numpy as _np
+        if isinstance(frames_raw, _np.ndarray) and frames_raw.ndim == 5:
+            frames_raw = frames_raw[0]
+    except Exception:
+        pass
+    return list(frames_raw)
+
+
 def _run_pipeline(pipe, kw: dict):
     result = pipe(**kw)
     # NB: `getattr(result, 'frames', None) or result[0]` is WRONG — when .frames
@@ -2449,7 +2605,21 @@ def _generate_video(pipe, request: VideoGenerationRequest):
     # Left loaded after the run so the next clip with the same set pays nothing.
     _sync_video_loras(pipe, getattr(request, 'loras', None))
     try:
-        frames = _run_pipeline(pipe, kw)
+        # Parts of this pipeline on other machines (codai/cluster/components.py),
+        # or this machine doing ITS part of someone else's run.
+        _handoff = getattr(request, '_handoff', None)
+        if isinstance(_handoff, dict) and _handoff.get('op'):
+            frames = _run_handoff(pipe, kw, request, _handoff)
+        else:
+            from codai.cluster.components import components_for as _components_for
+            _comp = _components_for(getattr(request, 'model', '') or '')
+            if _comp and type(pipe).__name__ in ('WanPipeline', 'WanImageToVideoPipeline'):
+                frames = _run_pipeline_split(pipe, kw, request, _comp)
+            else:
+                if _comp:
+                    print(f"  [video][components] {type(pipe).__name__} cannot be split "
+                          f"across machines; running here", flush=True)
+                frames = _run_pipeline(pipe, kw)
     except TaskCancelled:
         _vid_progress_done()
         raise  # global handler finishes the task (cancelled) + returns HTTP 499
@@ -4406,6 +4576,11 @@ async def video_generations(request: VideoGenerationRequest,
         raise HTTPException(status_code=500,
                             detail=f"Video generation failed: {_last_err}")
 
+    # A component run (or a finished video from the machine holding the VAE):
+    # the answer is already what the caller wants.
+    if isinstance(frames, _HandoffPayload):
+        return JSONResponse(frames.payload)
+
     # Encode raw frames to MP4 (per-model output quality via CRF when configured).
     try:
         import imageio, numpy as np

codai/cluster/components.py

+# CoderAI - OpenAI-compatible API server
+# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+
+"""A diffusion pipeline's parts on different machines.
+
+A video model is several models: a text encoder (UMT5-XXL, 11 GB in fp16),
+one or two denoising experts (Wan 2.2: high-noise then low-noise, ~10 GB
+each at 4-bit) and a VAE. They run one after the other, and what passes
+between them is small — prompt embeddings, a latent tensor of a few MB —
+so each can live on a different machine and the run is a relay over a LAN,
+not a tensor-parallel all-reduce. Per model, ``components`` in models.json:
+
+    {"text_encoder": "box2", "low_noise": "box3", "vae": "box3"}
+
+Each value is a cluster node name (or a URL block). The machine that took
+the request keeps what is not named: it encodes the prompt unless
+``text_encoder`` says otherwise, runs the high-noise expert (or the whole
+loop when the model has one expert), hands the latents to ``low_noise``'s
+machine at the expert boundary, and decodes them unless ``vae`` says
+otherwise — in which case the far machine returns the finished video.
+
+The far side runs the ordinary /v1/video/generations with a ``_handoff``
+block: the same request, so model loading, acceleration and LoRAs happen
+exactly as for a normal generation there; only what the loop starts from
+and what it returns differ (codai/api/video.py, _generate_video).
+"""
+
+import base64
+import io
+from typing import Optional
+
+_KEYS = ("text_encoder", "low_noise", "vae")
+
+
+def parse_components(raw) -> dict:
+    """{"text_encoder": node, "low_noise": node, "vae": node} — blank = here."""
+    d = raw if isinstance(raw, dict) else {}
+    out = {}
+    for k in _KEYS:
+        v = d.get(k)
+        if isinstance(v, str) and v.strip() and v.strip().lower() not in ("local", "here", "auto"):
+            out[k] = v.strip()
+        elif isinstance(v, dict) and v.get("url"):
+            out[k] = dict(v)
+    return out
+
+
+def components_for(model_name: str) -> dict:
+    """This model's placement, from its models.json entry."""
+    try:
+        from codai.models.manager import _model_entry_for
+        entry = _model_entry_for(model_name) or {}
+    except Exception:
+        return {}
+    return parse_components(entry.get("components"))
+
+
+# ------------------------------------------------------------- tensors
+def tensor_to_b64(t) -> str:
+    """A tensor as safetensors bytes, base64 — dtype and shape travel with it."""
+    from safetensors.torch import save
+    return base64.b64encode(save({"t": t.detach().contiguous().cpu()})).decode()
+
+
+def b64_to_tensor(s: str, device=None, dtype=None):
+    from safetensors.torch import load
+    t = load(base64.b64decode(s))["t"]
+    if dtype is not None:
+        t = t.to(dtype)
+    if device is not None:
+        t = t.to(device)
+    return t
+
+
+# --------------------------------------------------------------- remote
+def _peer(spec) -> dict:
+    """A node name or a {url, api_key} block → {name, url, api_key, verify}."""
+    if isinstance(spec, dict):
+        return {"name": spec.get("name") or spec.get("url"), "url": str(spec["url"]).rstrip("/"),
+                "api_key": spec.get("api_key", ""), "verify": spec.get("verify", "system")}
+    from codai.cluster.ddp import resolve_peers
+    return resolve_peers([{"name": spec}])[0]
+
+
+def call_handoff(spec, request_json: dict, handoff: dict, timeout: float = 3600.0) -> dict:
+    """POST the same generation request with a ``_handoff`` block to the
+    machine that owns a component; returns the JSON it answers."""
+    from codai.api import pod_http
+    peer = _peer(spec)
+    body = dict(request_json)
+    body["_handoff"] = handoff
+    headers = {"Content-Type": "application/json"}
+    if peer.get("api_key"):
+        headers["Authorization"] = f"Bearer {peer['api_key']}"
+    r = pod_http.post(peer["url"] + "/v1/video/generations", json=body, headers=headers,
+                      timeout=timeout, verify=(False if peer.get("verify") == "off" else True))
+    if r.status_code != 200:
+        raise RuntimeError(f"{peer['name']}: {r.status_code} {(r.text or '')[:300]}")
+    return r.json()
+
+
+def request_json(request) -> dict:
+    """The generation request as the far side must receive it — every field
+    the client sent, minus anything a client never sends."""
+    try:
+        d = request.model_dump(exclude_none=True)
+    except Exception:
+        d = dict(getattr(request, "__dict__", {}))
+    d.pop("_handoff", None)
+    return d
+
+
+# ---------------------------------------------------- the wan step split
+def boundary_step(pipe, kw: dict) -> Optional[int]:
+    """The first loop index that uses the low-noise expert, or None when the
+    model has a single expert. Computed from the scheduler exactly as the
+    pipeline will (set_timesteps + the boundary ratio)."""
+    ratio = getattr(getattr(pipe, "config", None), "boundary_ratio", None)
+    if ratio is None or getattr(pipe, "transformer_2", None) is None:
+        return None
+    import copy
+    sched = copy.deepcopy(pipe.scheduler)
+    sched.set_timesteps(int(kw.get("num_inference_steps") or 25))
+    boundary = float(ratio) * float(sched.config.num_train_timesteps)
+    for i, t in enumerate(sched.timesteps):
+        if float(t) < boundary:
+            return i
+    return None
+
+
+class SlicedTimesteps:
+    """Make the pipeline's scheduler start at loop index ``start``: after the
+    pipeline's own ``set_timesteps`` the timestep and sigma tables are cut
+    so the loop runs only the remaining steps and ``step()`` still finds
+    each timestep's sigma pair."""
+
+    def __init__(self, scheduler, start: int):
+        self.scheduler = scheduler
+        self.start = int(start)
+        self._orig = None
+
+    def __enter__(self):
+        sched, start = self.scheduler, self.start
+        self._orig = sched.set_timesteps
+
+        def _set(*a, **k):
+            self._orig(*a, **k)
+            sched.timesteps = sched.timesteps[start:]
+            if getattr(sched, "sigmas", None) is not None:
+                sched.sigmas = sched.sigmas[start:]
+            if hasattr(sched, "_step_index"):
+                sched._step_index = None
+            if hasattr(sched, "_begin_index"):
+                sched._begin_index = None
+        sched.set_timesteps = _set
+        return self
+
+    def __exit__(self, *exc):
+        self.scheduler.set_timesteps = self._orig
+        return False