FLUX.2-klein PR 1: Scaffold for flow matching shared by several models, first of which is flux2-klein

docs/adding_models.rst

@@ -1641,6 +1641,74 @@ context segments have a non-zero span in the prefill step that writes them, and
 0 in every later step. The ``commit`` in the prefill step converts the reservation into
 resident pages that the later steps read.
 
+Worked example: an image DiT on the diffusion scaffold
+------------------------------------------------------
+
+Flow-matching image and video transformers (FLUX-family DiTs, Z-Image, Wan-style
+video) share one graph shape, so M* ships it as a scaffold in
+``mstar/model/components/diffusion/``; a model adds only its exact component
+ports, its checkpoint remap and a few hooks.
+
+Graph and Walks::
+
+    text_encoder ──text_embeds──▶ Loop("denoise_loop", dit) ──latents──▶ vae_decoder ──▶ EMIT image
+    vae_encoder  ──ref_latents──▶ (image_edit only)
+
+    encode_text  : text_inputs  -> text_embeds  (persist)
+    encode_image : image_inputs -> ref_latents  (persist)
+    image_gen    : Sequential[Loop(dit) -> vae_decoder]
+    image_edit   : image_gen with ref_latents as an extra Loop input
+
+``latents`` is the loop-back edge; every other input is re-injected each
+iteration. The step index is the engine's loop counter
+(``fwd_info.dynamic_loop_iter_counts``), so there is no host sync per step.
+
+The ``dit`` node subclasses ``DenoiseLoopSubmodule`` (``denoise_loop.py``), which owns
+the engine contract — seeded initial noise, per-request schedules, equal-key request
+batching (``can_batch``), stacked ``forward_batched``, ``check_stop`` at the request's own
+step count, the ragged-attention ``declare_step`` and one CUDA-graph bucket per
+``bucket_key`` with the Euler update inside the graph — and asks the model for:
+
+.. code-block:: python
+
+    class MyDenoise(DenoiseLoopSubmodule):
+        def bucket_key_for(self, fwd_info): ...       # hashable: latent grid, text length, ...
+        def schedule_for(self, fwd_info, key): ...    # FlowMatchSchedule (flow_match.py)
+        def seed_latents(self, fwd_info, key, gen): ...  # [L, C] initial noise, seeded
+        def request_inputs(self, fwd_info, inputs, key): ...  # per-row conditioning tensors
+        def num_tokens(self, key): ...                # image tokens for the attention step
+        def capture_request_inputs(self, key, device): ...   # dummy rows for graph capture
+        def denoise(self, engine_inputs, key, latents, timestep, sigma, sigma_next, **cond): ...
+        # optional: extra attention spans (a refiner over the image tokens alone, ...)
+        def attention_segments(self, key): return (("main", self.num_tokens(key)),)
+
+``bucket_key`` is whatever has to match for two rows to share a forward, and it is also
+the CUDA-graph bucket key — shape, but equally a knob like whether CFG is on.
+
+A multi-step scheduler that carries solver state between iterations (wan22's UniPC, with
+``unipc_model_outputs`` and ``unipc_last_sample``) names it in ``SOLVER_STATE``; those
+become loop-back edges alongside ``latents``, seeded by ``seed_loop_back`` and returned
+from ``denoise`` as a ``{name: tensor}`` mapping. A scheduler needing another per-step
+scalar adds a ``STEP_SCALARS`` entry naming the attribute of its ``FlowMatchSchedule``
+to slice it from — the base stages that tensor on the device and slices it per step, so
+nothing in the batching path has to change.
+
+Every span a layer attends over must be declared with its own label (default: one
+``"main"`` segment per request); ``ragged_for(label)`` hands that span's kernel to the
+layers, which fall back to SDPA when no ragged resource is bound.
+
+Shared pieces to build the components from: ``flow_match.py`` (shift schedules and the
+exact Euler step), ``rope.py`` (multi-axis rotary embeddings), ``attention.py`` (joint
+attention over packed tokens: SDPA eagerly, the ragged FlashInfer resource under graphs),
+``text_encoder.py`` (native Qwen3 hidden-state taps), ``autoencoder_kl.py`` (the KL VAE
+family), ``image_io.py`` (latent packing, PNG) and ``lora.py`` (static LoRA merge at load,
+also into fused ``q``/``k``/``v`` shards via ``FusedColumnLinear.shard_slice``).
+
+On the API side ``DiffusionImageAdapter`` (``mstar/api_server/openai/adapters.py``) serves
+``/v1/images/generations`` and ``/v1/images/edits`` for any such model (``size`` becomes
+``width``/``height``, ``seed`` and other knobs pass through); register it under the
+model's registry key. The SDK exposes ``generate_image`` and ``edit_image``.
+
 Advanced: async partitions and streaming
 ----------------------------------------
 

docs/clients.rst

@@ -151,6 +151,9 @@ Convenience wrappers:
      - Text generation (and, with ``output_modalities=("text", "audio")``, speech).
    * - ``generate_image(prompt, **kw)``
      - PNG ``bytes`` (e.g. BAGEL text-to-image).
+   * - ``edit_image(prompt, image, **kw)``
+     - PNG ``bytes`` for an edit of ``image`` (a path, bytes, or a list of them for
+       multi-reference models).
    * - ``tts(text, *, voice=None, **kw)``
      - An ``AudioBuffer`` (``.to_wav(path)``, ``.to_numpy()``, ``len(...)`` samples).
    * - ``voices()``

examples/sdk_image_edit.py

@@ -0,0 +1,18 @@
+"""Text-to-image, then an edit of the result, via the Python SDK.
+
+Start a server for an image model that serves edits first, e.g.:  mstar serve bagel
+"""
+
+from mstar import MStarClient
+
+client = MStarClient("http://localhost:8000")
+
+png = client.generate_image("a cat holding a sign that says hello world", seed=0)
+with open("out.png", "wb") as f:
+    f.write(png)
+print(f"wrote out.png - {len(png)} bytes")
+
+edited = client.edit_image("make the sign say goodbye, watercolor style", "out.png", seed=1)
+with open("edit.png", "wb") as f:
+    f.write(edited)
+print(f"wrote edit.png - {len(edited)} bytes")

mstar/api_server/openai/adapters.py

@@ -645,6 +645,53 @@ def video_to_request(self, req: VideoGenerationRequest, upload_dir: Path) -> Sub
         )
 
 
+class DiffusionImageAdapter(OpenAIAdapter):
+    """Generic adapter for the DiT-scaffold image models (FLUX.2 klein, Z-Image, ...):
+    text-to-image and reference-image editing.
+
+    ``size`` ("WxH") maps to the models' ``width`` / ``height`` kwargs; ``seed`` is
+    honoured by the conductor; every other knob (``num_inference_steps``,
+    ``guidance_scale`` on models that take it, ``negative_prompt``) passes through
+    ``extra_body``. ``n`` is served as ``n`` concurrent requests by the images handler,
+    which the equal-shape batching then runs together.
+    """
+
+    supports_images = True
+
+    @staticmethod
+    def _size_kwargs(size: str | None, mk: dict) -> None:
+        if not size:
+            return
+        try:
+            width, height = (int(v) for v in size.lower().split("x"))
+        except ValueError:
+            raise ValueError(f"size must be 'WxH' (e.g. '1024x1024'); got {size!r}") from None
+        mk.setdefault("width", width)
+        mk.setdefault("height", height)
+
+    def image_to_request(self, req: ImageGenerationRequest, upload_dir: Path) -> SubmitArgs:  # noqa: ARG002
+        mk = _passthrough(req)
+        self._size_kwargs(getattr(req, "size", None), mk)
+        if getattr(req, "seed", None) is not None:
+            mk.setdefault("seed", req.seed)
+        return SubmitArgs(
+            text=req.prompt,
+            input_modalities=["text"],
+            output_modalities=["image"],
+            model_kwargs=mk,
+        )
+
+    def image_edit_to_request(self, prompt: str, image_path: str, extra_kwargs: dict) -> SubmitArgs:
+        mk = dict(extra_kwargs or {})
+        self._size_kwargs(mk.pop("size", None), mk)
+        return SubmitArgs(
+            text=prompt,
+            file_paths={"image": [image_path]},
+            input_modalities=["image", "text"],
+            output_modalities=["image"],
+            model_kwargs=mk,
+        )
+
 # Whisper-style control tokens: ``<|en|>`` (language), ``<|12.34|>`` (timestamp),
 # and the task/format markers. ASR models render the ones that carry
 # information (language, timestamps) into their text stream so the adapter can

mstar/api_server/openai/router.py

@@ -191,7 +191,9 @@ async def images_generations(request: ImageGenerationRequest, raw_request: Reque
     try:
         result = await serving_images.create_images(api, model_name, adapter, request, raw_request)
     except Exception as e:  # noqa: BLE001
-        return _error(getattr(e, "status_code", 500), str(getattr(e, "detail", e)), "server_error")
+        # a malformed request (adapter / model validation) is the client's error, not the server's
+        default_status = 400 if isinstance(e, (ValueError, TypeError)) else 500
+        return _error(getattr(e, "status_code", default_status), str(getattr(e, "detail", e)), "server_error")
     return JSONResponse(result)
 
 
@@ -237,6 +239,8 @@ async def images_edits(request: Request):
                 extra[key] = json.loads(value)
             except (json.JSONDecodeError, TypeError):
                 extra[key] = value
+        if form.get("size"):
+            extra["size"] = form.get("size")  # the adapter maps "WxH" to the model's width / height
         result = await serving_images.create_image_edit(
             api,
             model_name,
@@ -248,5 +252,6 @@ async def images_edits(request: Request):
             raw_request=request,
         )
     except Exception as e:  # noqa: BLE001
-        return _error(getattr(e, "status_code", 500), str(getattr(e, "detail", e)), "server_error")
+        default_status = 400 if isinstance(e, (ValueError, TypeError)) else 500
+        return _error(getattr(e, "status_code", default_status), str(getattr(e, "detail", e)), "server_error")
     return JSONResponse(result)

mstar/client/client.py

@@ -11,6 +11,7 @@
     print(client.chat("Hello!").text)
     client.tts("Hi there", voice="tara").to_wav("out.wav")
     open("cat.png", "wb").write(client.generate_image("a cat in a hat"))
+    open("edit.png", "wb").write(client.edit_image("make it a watercolor", "cat.png"))
 """
 
 from __future__ import annotations
@@ -178,6 +179,17 @@ def generate_image(self, prompt: str, **model_kwargs) -> bytes:
             raise RuntimeError("Server returned no image output")
         return res.images[0]
 
+    def edit_image(self, prompt: str, image, **model_kwargs) -> bytes:
+        """Return PNG bytes for an image-editing request: ``image`` (a path, raw
+        bytes or a ``(filename, bytes)`` tuple, or a list of them for multi-reference
+        models) plus the edit instruction (e.g. BAGEL)."""
+        res = self.generate(
+            text=prompt, images=image, output_modalities=("image",), **model_kwargs,
+        )
+        if not res.images:
+            raise RuntimeError("Server returned no image output")
+        return res.images[0]
+
     def tts(
         self,
         text: str,
@@ -256,6 +268,7 @@ def _build_files(self, images, audio, video) -> list[tuple[str, tuple[str, bytes
                 and isinstance(items[0], str)
                 and isinstance(items[1], (bytes, bytearray))
             )
+            # one path, one blob, or one (filename, bytes) pair
             if isinstance(items, (str, bytes, bytearray, Path)) or named_bytes:
                 items = [items]
             for i, item in enumerate(items):

mstar/engine/resources/convenience.py

@@ -1,6 +1,7 @@
 import torch
 
 from mstar.engine.resources.attn.base import AttentionManager
+from mstar.engine.resources.attn.ragged.base import RaggedAttnManager
 from mstar.engine.resources.kv.manager import KVManager
 
 
@@ -18,6 +19,8 @@ class AttentionCallable:
     That sharing is why the label is one cursor for the whole stack. No model
     varies its label per layer today; one that needs to should thread the label
     explicitly rather than use this.
+
+    For a cacheless attention, see :class:`RaggedAttentionCallable`.
     """
 
     def __init__(self, kv: KVManager, attn: AttentionManager | None=None):
@@ -56,3 +59,32 @@ def __call__(
         if self.attn.requires_kv_write:
             self.kv.write_kv(k, v)
         return self.attn.run(q, kv_cache_layer=self.kv.layer_view(), k=k, v=v)
+
+
+class RaggedAttentionCallable:
+    """`AttentionCallable` for a cacheless attention: ``(q, k, v) -> out`` over
+    one declared span of a `RaggedAttnManager`.
+
+    No ``kv`` and no layer cursor, because there is neither to carry: a ragged
+    plan is keyed by the segment label, and the wrapper behind it is per
+    (bucket, slot, label) rather than per layer. The label is therefore bound
+    here at construction instead of being a cursor on the resource -- which is
+    what lets one stack attend over several spans at once (a refiner over the
+    image tokens next to a pass over all of them), the case `AttentionCallable`
+    says to thread explicitly.
+
+    One instance per label, held for the life of the resource binding: a
+    compiled transformer region guards on the identity of the callables it is
+    handed, so handing it a fresh one per step retraces that frame every step
+    until Dynamo gives up.
+    """
+
+    def __init__(self, attn: RaggedAttnManager, label: str):
+        self.attn = attn
+        self.label = label
+
+    @torch.compiler.disable
+    def __call__(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.attn.run(q, k, v, label=self.label)

mstar/graph/base.py

@@ -324,6 +324,16 @@ def get_inputs_outputs(self):
             loop_back=loop_back
         )
 
+    def persisted_input_names(self) -> set[str]:
+        """Inputs the enclosing loop re-injects unchanged into ``ready_signals``
+        every iteration (``Loop.ingest_external_input`` / ``complete_iter``): a
+        denoise step's text conditioning, say. They never land in
+        ``ready_next_iter``, but they are as good as ready for the next one."""
+        return {
+            name for name, edge in self.ready_signals.ready_inputs.items()
+            if edge._persist_for_loop
+        }
+
     def is_ready_for_speculation(
         self, check_next_iter: bool=False,
         allow_streaming: bool=True
@@ -332,19 +342,12 @@ def is_ready_for_speculation(
         if allow_streaming:
             needed_inputs = needed_inputs - self._streaming_inputs
         if check_next_iter:
-            # Loop-external inputs are re-injected unchanged into ready_signals
-            # every iteration (Loop.ingest_external_input / complete_iter) and
-            # never routed through ready_next_iter. Count them here so a
-            # same-node next-iteration speculation isn't blocked on inputs
-            # that are guaranteed to reappear.
-            carried_names = {
-                name for name, edge in self.ready_signals.ready_inputs.items()
-                if edge._persist_for_loop
-            }
+            # Count the persisted inputs here so a same-node next-iteration
+            # speculation isn't blocked on inputs guaranteed to reappear.
             return needed_inputs.issubset(
                 self.ready_next_iter.ready_names
                 | self.speculative_signals.ready_names
-                | carried_names
+                | self.persisted_input_names()
             )
         return needed_inputs.issubset(
             self.ready_signals.ready_names | self.speculative_signals.ready_names

mstar/graph/runtime/python.py

@@ -898,12 +898,10 @@ def _prep_one_spec_rid(
             else node.ready_signals.ready_inputs
         )
         if same_node:
-            # Carry over loop-external inputs sitting in ready_signals (see
-            # GraphNode.is_ready_for_speculation): they are re-injected
-            # unchanged every iteration and never land in ready_next_iter.
-            for name, edge in node.ready_signals.ready_inputs.items():
-                if edge._persist_for_loop and name not in inputs:
-                    inputs[name] = edge
+            # Carry over the loop-external inputs sitting in ready_signals —
+            # the same set `is_ready_for_speculation` just counted as ready.
+            for name in node.persisted_input_names():
+                inputs.setdefault(name, node.ready_signals.ready_inputs[name])
         return _SpecRidPrep(
             rid=rid,
             node=node,

mstar/model/components/diffusion/__init__.py

@@ -0,0 +1,16 @@
+"""The DiT scaffold: shared pieces of a text -> flow-transformer loop -> VAE model.
+
+Each module here is model-agnostic and tested against the reference library
+once, so a new image/video diffusion model brings only its config, its exact
+component ports, a weight loader and a thin graph file:
+
+* ``flow_match``   sigma schedules (linear / dynamic / FLUX.2 empirical shift) + the Euler step
+* ``rope``         multi-axis rotary tables and the interleaved application FLUX-family DiTs use
+* ``attention``    joint attention over packed text+image tokens: SDPA eagerly, the ragged
+                   FlashInfer resource under CUDA graphs
+* ``text_encoder`` a native Qwen3-family encoder returning tapped hidden states
+* ``denoise_loop`` the Loop-body submodule: seeded noise, equal-shape request batching, per-step
+                   CUDA graphs, stop at the request's step count
+* ``image_io``     pixel/latent packing helpers and PNG encoding
+* ``lora``         static LoRA merging at load time
+"""

mstar/model/components/diffusion/attention.py

@@ -0,0 +1,53 @@
+"""Joint attention over a DiT's packed tokens.
+
+A diffusion transformer attends bidirectionally over all of a request's tokens
+(text + image + reference-image) and keeps nothing between steps, so there is
+no KV cache to page. Two backends serve the same ``[B, S, H, D]`` contract:
+
+* ``sdpa``   ``torch.nn.functional.scaled_dot_product_attention`` — the reference
+             pipelines' own kernel, hence the parity baseline; used eagerly and on CPU.
+* the engine's **ragged attention resource** (``RaggedAttentionSpec``, FlashInfer's
+  ragged prefill on FA2/FA3) — one segment per request, planned by the runner
+  outside the captured graph, so a per-step CUDA graph can replay it. The
+  resource is bound by the submodule and threaded to the layers as a callable.
+
+Layers call :func:`joint_attention`; the choice is made once per forward by
+whoever owns the resource, not per layer.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+import torch.nn.functional as F
+
+# ``(q, k, v) -> out`` over packed ``[total_tokens, H, D]`` tensors, one segment per request.
+RaggedAttentionFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
+
+
+def sdpa_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """Bidirectional attention over ``[B, S, H, D]`` tensors via SDPA; returns ``[B, S, H, D]``."""
+    out = F.scaled_dot_product_attention(
+        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=False,
+    )
+    return out.transpose(1, 2)
+
+
+def joint_attention(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ragged: RaggedAttentionFn | None = None,
+) -> torch.Tensor:
+    """Attention over ``[B, S, H, D]`` q/k/v (all requests in the batch share
+    ``S``) through the ragged resource when one is bound, else SDPA.
+
+    The ragged kernel wants the batch packed to ``[B * S, H, D]`` with one
+    segment of ``S`` tokens per request — exactly what the owning submodule
+    declared in its step — and hands back the same packing.
+    """
+    if ragged is None:
+        return sdpa_attention(q, k, v)
+    bsz, seq, heads, dim = q.shape
+    out = ragged(
+        q.reshape(bsz * seq, heads, dim), k.reshape(bsz * seq, heads, dim), v.reshape(bsz * seq, heads, dim),
+    )
+    return out.view(bsz, seq, heads, dim).to(q.dtype)