state compaction

introduces a new topic ("state_catchup") over which a full state can be sent. currently the master sends the worker + api this new state, and they update only if they have no other events applied - otherwise usual NACK systems function ## testing manually tested on two nodes
Enable UI settings for image editing (#1258 )
2026-01-23 13:29:29 -05:00 · 2026-01-23 14:15:06 +00:00 · 2026-01-23 13:37:25 +00:00 · 2026-01-23 11:19:58 +00:00
12 changed files with 446 additions and 185 deletions
--- a/dashboard/src/lib/components/ChatForm.svelte
+++ b/dashboard/src/lib/components/ChatForm.svelte
@@ -89,7 +89,10 @@

  const isImageModel = $derived(() => {
    if (!currentModel) return false;
-    return modelSupportsTextToImage(currentModel);
+    return (
+      modelSupportsTextToImage(currentModel) ||
+      modelSupportsImageEditing(currentModel)
+    );
  });

  const isEditOnlyWithoutImage = $derived(
@@ -646,6 +649,23 @@
            </svg>
            <span>EDIT</span>
          </span>
+        {:else if isEditOnlyWithoutImage}
+          <span class="inline-flex items-center gap-1.5">
+            <svg
+              class="w-3.5 h-3.5"
+              fill="none"
+              viewBox="0 0 24 24"
+              stroke="currentColor"
+              stroke-width="2"
+            >
+              <path
+                stroke-linecap="round"
+                stroke-linejoin="round"
+                d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z"
+              />
+            </svg>
+            <span>EDIT</span>
+          </span>
        {:else if isImageModel()}
          <span class="inline-flex items-center gap-1.5">
            <svg
--- a/dashboard/src/lib/components/ImageParamsPanel.svelte
+++ b/dashboard/src/lib/components/ImageParamsPanel.svelte
@@ -110,6 +110,36 @@
    setImageGenerationParams({ negativePrompt: value || null });
  }

+  function handleNumImagesChange(event: Event) {
+    const input = event.target as HTMLInputElement;
+    const value = input.value.trim();
+    if (value === "") {
+      setImageGenerationParams({ numImages: 1 });
+    } else {
+      const num = parseInt(value, 10);
+      if (!isNaN(num) && num >= 1) {
+        setImageGenerationParams({ numImages: num });
+      }
+    }
+  }
+
+  function handleStreamChange(enabled: boolean) {
+    setImageGenerationParams({ stream: enabled });
+  }
+
+  function handlePartialImagesChange(event: Event) {
+    const input = event.target as HTMLInputElement;
+    const value = input.value.trim();
+    if (value === "") {
+      setImageGenerationParams({ partialImages: 0 });
+    } else {
+      const num = parseInt(value, 10);
+      if (!isNaN(num) && num >= 0) {
+        setImageGenerationParams({ partialImages: num });
+      }
+    }
+  }
+
  function clearSteps() {
    setImageGenerationParams({ numInferenceSteps: null });
  }
@@ -134,90 +164,92 @@
 <div class="border-b border-exo-medium-gray/30 px-3 py-2">
  <!-- Basic params row -->
  <div class="flex items-center gap-3 flex-wrap">
-    <!-- Size -->
-    <div class="flex items-center gap-1.5">
-      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
-        >SIZE:</span
-      >
-      <div class="relative">
-        <button
-          bind:this={sizeButtonRef}
-          type="button"
-          onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
-          class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
-            ? 'border-exo-yellow/70'
-            : ''}"
+    <!-- Size (hidden in edit mode - output size comes from input image) -->
+    {#if !isEditMode}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >SIZE:</span
        >
-          {params.size}
-        </button>
-        <div
-          class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
-            ? 'rotate-180'
-            : ''}"
-        >
-          <svg
-            class="w-3 h-3 text-exo-yellow/60"
-            fill="none"
-            viewBox="0 0 24 24"
-            stroke="currentColor"
+        <div class="relative">
+          <button
+            bind:this={sizeButtonRef}
+            type="button"
+            onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
+            class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
+              ? 'border-exo-yellow/70'
+              : ''}"
          >
-            <path
-              stroke-linecap="round"
-              stroke-linejoin="round"
-              stroke-width="2"
-              d="M19 9l-7 7-7-7"
-            />
-          </svg>
-        </div>
-      </div>
-
-      {#if isSizeDropdownOpen}
-        <!-- Backdrop to close dropdown -->
-        <button
-          type="button"
-          class="fixed inset-0 z-[9998] cursor-default"
-          onclick={() => (isSizeDropdownOpen = false)}
-          aria-label="Close dropdown"
-        ></button>
-
-        <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
-        <div
-          class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto min-w-max"
-          style="bottom: calc(100vh - {sizeDropdownPosition()
-            .top}px + 4px); left: {sizeDropdownPosition().left}px;"
-        >
-          <div class="py-1">
-            {#each sizeOptions as size}
-              <button
-                type="button"
-                onclick={() => selectSize(size)}
-                class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
-                size
-                  ? 'bg-transparent text-exo-yellow'
-                  : 'text-exo-light-gray hover:text-exo-yellow'}"
-              >
-                {#if params.size === size}
-                  <svg
-                    class="w-3 h-3 flex-shrink-0"
-                    fill="currentColor"
-                    viewBox="0 0 20 20"
-                  >
-                    <path
-                      fill-rule="evenodd"
-                      d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
-                      clip-rule="evenodd"
-                    />
-                  </svg>
-                {:else}
-                  <span class="w-3"></span>
-                {/if}
-                <span>{size}</span>
-              </button>
-            {/each}
+            {params.size}
+          </button>
+          <div
+            class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
+              ? 'rotate-180'
+              : ''}"
+          >
+            <svg
+              class="w-3 h-3 text-exo-yellow/60"
+              fill="none"
+              viewBox="0 0 24 24"
+              stroke="currentColor"
+            >
+              <path
+                stroke-linecap="round"
+                stroke-linejoin="round"
+                stroke-width="2"
+                d="M19 9l-7 7-7-7"
+              />
+            </svg>
          </div>
        </div>
-      {/if}
-    </div>
+
+        {#if isSizeDropdownOpen}
+          <!-- Backdrop to close dropdown -->
+          <button
+            type="button"
+            class="fixed inset-0 z-[9998] cursor-default"
+            onclick={() => (isSizeDropdownOpen = false)}
+            aria-label="Close dropdown"
+          ></button>
+
+          <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
+          <div
+            class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto min-w-max"
+            style="bottom: calc(100vh - {sizeDropdownPosition()
+              .top}px + 4px); left: {sizeDropdownPosition().left}px;"
+          >
+            <div class="py-1">
+              {#each sizeOptions as size}
+                <button
+                  type="button"
+                  onclick={() => selectSize(size)}
+                  class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
+                  size
+                    ? 'bg-transparent text-exo-yellow'
+                    : 'text-exo-light-gray hover:text-exo-yellow'}"
+                >
+                  {#if params.size === size}
+                    <svg
+                      class="w-3 h-3 flex-shrink-0"
+                      fill="currentColor"
+                      viewBox="0 0 20 20"
+                    >
+                      <path
+                        fill-rule="evenodd"
+                        d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
+                        clip-rule="evenodd"
+                      />
+                    </svg>
+                  {:else}
+                    <span class="w-3"></span>
+                  {/if}
+                  <span>{size}</span>
+                </button>
+              {/each}
+            </div>
+          </div>
+        {/if}
+      </div>
+    {/if}

    <!-- Quality -->
    <div class="flex items-center gap-1.5">
@@ -325,6 +357,59 @@
      </div>
    </div>

+    <!-- Number of Images (not in edit mode) -->
+    {#if !isEditMode}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >IMAGES:</span
+        >
+        <input
+          type="number"
+          min="1"
+          value={params.numImages}
+          oninput={handleNumImagesChange}
+          class="w-12 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow text-center transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+        />
+      </div>
+    {/if}
+
+    <!-- Stream toggle -->
+    <div class="flex items-center gap-1.5">
+      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+        >STREAM:</span
+      >
+      <button
+        type="button"
+        onclick={() => handleStreamChange(!params.stream)}
+        class="w-8 h-4 rounded-full transition-all duration-200 cursor-pointer relative {params.stream
+          ? 'bg-exo-yellow'
+          : 'bg-exo-medium-gray/50 border border-exo-yellow/30'}"
+        title={params.stream ? "Streaming enabled" : "Streaming disabled"}
+      >
+        <div
+          class="absolute top-0.5 w-3 h-3 rounded-full transition-all duration-200 {params.stream
+            ? 'right-0.5 bg-exo-black'
+            : 'left-0.5 bg-exo-light-gray'}"
+        ></div>
+      </button>
+    </div>
+
+    <!-- Partial Images (only when streaming) -->
+    {#if params.stream}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >PARTIALS:</span
+        >
+        <input
+          type="number"
+          min="0"
+          value={params.partialImages}
+          oninput={handlePartialImagesChange}
+          class="w-12 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow text-center transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+        />
+      </div>
+    {/if}
+
    <!-- Input Fidelity (edit mode only) -->
    {#if isEditMode}
      <div class="flex items-center gap-1.5">
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -238,6 +238,10 @@ export interface ImageGenerationParams {
  size: "512x512" | "768x768" | "1024x1024" | "1024x768" | "768x1024";
  quality: "low" | "medium" | "high";
  outputFormat: "png" | "jpeg";
+  numImages: number;
+  // Streaming params
+  stream: boolean;
+  partialImages: number;
  // Advanced params
  seed: number | null;
  numInferenceSteps: number | null;
@@ -257,6 +261,9 @@ const DEFAULT_IMAGE_PARAMS: ImageGenerationParams = {
  size: "1024x1024",
  quality: "medium",
  outputFormat: "png",
+  numImages: 1,
+  stream: true,
+  partialImages: 3,
  seed: null,
  numInferenceSteps: null,
  guidance: null,
@@ -1809,12 +1816,13 @@ class AppStore {
      const requestBody: Record<string, unknown> = {
        model,
        prompt,
+        n: params.numImages,
        quality: params.quality,
        size: params.size,
        output_format: params.outputFormat,
        response_format: "b64_json",
-        stream: true,
-        partial_images: 3,
+        stream: params.stream,
+        partial_images: params.partialImages,
      };

      if (hasAdvancedParams) {
@@ -1878,31 +1886,74 @@ class AppStore {
              if (imageData && idx !== -1) {
                const format = parsed.format || "png";
                const mimeType = `image/${format}`;
+                const imageIndex = parsed.image_index ?? 0;
+                const numImages = params.numImages;
+
                if (parsed.type === "partial") {
                  // Update with partial image and progress
                  const partialNum = (parsed.partial_index ?? 0) + 1;
                  const totalPartials = parsed.total_partials ?? 3;
-                  this.messages[idx].content =
-                    `Generating... ${partialNum}/${totalPartials}`;
-                  this.messages[idx].attachments = [
-                    {
-                      type: "generated-image",
-                      name: `generated-image.${format}`,
-                      preview: `data:${mimeType};base64,${imageData}`,
-                      mimeType,
-                    },
-                  ];
+                  const progressText =
+                    numImages > 1
+                      ? `Generating image ${imageIndex + 1}/${numImages}... ${partialNum}/${totalPartials}`
+                      : `Generating... ${partialNum}/${totalPartials}`;
+                  this.messages[idx].content = progressText;
+
+                  const partialAttachment: MessageAttachment = {
+                    type: "generated-image",
+                    name: `generated-image.${format}`,
+                    preview: `data:${mimeType};base64,${imageData}`,
+                    mimeType,
+                  };
+
+                  if (imageIndex === 0) {
+                    // First image - safe to replace attachments with partial preview
+                    this.messages[idx].attachments = [partialAttachment];
+                  } else {
+                    // Subsequent images - keep existing finals, show partial at current position
+                    const existingAttachments =
+                      this.messages[idx].attachments || [];
+                    // Keep only the completed final images (up to current imageIndex)
+                    const finals = existingAttachments.slice(0, imageIndex);
+                    this.messages[idx].attachments = [
+                      ...finals,
+                      partialAttachment,
+                    ];
+                  }
                } else if (parsed.type === "final") {
-                  // Final image
-                  this.messages[idx].content = "";
-                  this.messages[idx].attachments = [
-                    {
-                      type: "generated-image",
-                      name: `generated-image.${format}`,
-                      preview: `data:${mimeType};base64,${imageData}`,
-                      mimeType,
-                    },
-                  ];
+                  // Final image - replace partial at this position
+                  const newAttachment: MessageAttachment = {
+                    type: "generated-image",
+                    name: `generated-image-${imageIndex + 1}.${format}`,
+                    preview: `data:${mimeType};base64,${imageData}`,
+                    mimeType,
+                  };
+
+                  if (imageIndex === 0) {
+                    // First final image - replace any partial preview
+                    this.messages[idx].attachments = [newAttachment];
+                  } else {
+                    // Subsequent images - keep previous finals, replace partial at current position
+                    const existingAttachments =
+                      this.messages[idx].attachments || [];
+                    // Slice keeps indices 0 to imageIndex-1 (the previous final images)
+                    const previousFinals = existingAttachments.slice(
+                      0,
+                      imageIndex,
+                    );
+                    this.messages[idx].attachments = [
+                      ...previousFinals,
+                      newAttachment,
+                    ];
+                  }
+
+                  // Update progress message for multiple images
+                  if (numImages > 1 && imageIndex < numImages - 1) {
+                    this.messages[idx].content =
+                      `Generating image ${imageIndex + 2}/${numImages}...`;
+                  } else {
+                    this.messages[idx].content = "";
+                  }
                }
              }
            } catch {
@@ -1983,8 +2034,8 @@ class AppStore {
      formData.append("size", params.size);
      formData.append("output_format", params.outputFormat);
      formData.append("response_format", "b64_json");
-      formData.append("stream", "1"); // Use "1" instead of "true" for reliable FastAPI boolean parsing
-      formData.append("partial_images", "3");
+      formData.append("stream", params.stream ? "1" : "0");
+      formData.append("partial_images", params.partialImages.toString());
      formData.append("input_fidelity", params.inputFidelity);

      // Advanced params
--- a/src/exo/main.py
+++ b/src/exo/main.py
@@ -49,6 +49,7 @@ class Node:
        await router.register_topic(topics.COMMANDS)
        await router.register_topic(topics.ELECTION_MESSAGES)
        await router.register_topic(topics.CONNECTION_MESSAGES)
+        await router.register_topic(topics.STATE_CATCHUP)

        logger.info(f"Starting node {node_id}")
        if args.spawn_api:
@@ -59,6 +60,7 @@ class Node:
                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                command_sender=router.sender(topics.COMMANDS),
                election_receiver=router.receiver(topics.ELECTION_MESSAGES),
+                state_catchup_receiver=router.receiver(topics.STATE_CATCHUP),
            )
        else:
            api = None
@@ -72,6 +74,7 @@ class Node:
                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                local_event_sender=router.sender(topics.LOCAL_EVENTS),
                command_sender=router.sender(topics.COMMANDS),
+                state_catchup_receiver=router.receiver(topics.STATE_CATCHUP),
            )
        else:
            worker = None
@@ -83,6 +86,7 @@ class Node:
            global_event_sender=router.sender(topics.GLOBAL_EVENTS),
            local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
            command_receiver=router.receiver(topics.COMMANDS),
+            state_catchup_sender=router.sender(topics.STATE_CATCHUP),
        )

        er_send, er_recv = channel[ElectionResult]()
@@ -153,6 +157,7 @@ class Node:
                        global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
                        local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
                        command_receiver=self.router.receiver(topics.COMMANDS),
+                        state_catchup_sender=self.router.sender(topics.STATE_CATCHUP),
                    )
                    self._tg.start_soon(self.master.run)
                elif (
@@ -185,6 +190,9 @@ class Node:
                            ),
                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                            command_sender=self.router.sender(topics.COMMANDS),
+                            state_catchup_receiver=self.router.receiver(
+                                topics.STATE_CATCHUP
+                            ),
                        )
                        self._tg.start_soon(self.worker.run)
                    if self.api:
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -1,4 +1,5 @@
 import base64
+import contextlib
 import json
 import time
 from collections.abc import AsyncGenerator
@@ -33,6 +34,7 @@ from exo.shared.models.model_cards import (
    ModelId,
 )
 from exo.shared.types.api import (
+    AdvancedImageParams,
    BenchChatCompletionResponse,
    BenchChatCompletionTaskParams,
    BenchImageGenerationResponse,
@@ -156,12 +158,14 @@ class API:
        command_sender: Sender[ForwarderCommand],
        # This lets us pause the API if an election is running
        election_receiver: Receiver[ElectionMessage],
+        state_catchup_receiver: Receiver[State],
    ) -> None:
        self.state = State()
        self._event_log: list[Event] = []
        self.command_sender = command_sender
        self.global_event_receiver = global_event_receiver
        self.election_receiver = election_receiver
+        self.state_catchup_receiver = state_catchup_receiver
        self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
        self.node_id: NodeId = node_id
        self.session_id: SessionId = session_id
@@ -835,6 +839,7 @@ class API:
                            # Yield partial image event (always use b64_json for partials)
                            event_data = {
                                "type": "partial",
+                                "image_index": chunk.image_index,
                                "partial_index": partial_idx,
                                "total_partials": total_partials,
                                "format": str(chunk.format),
@@ -1024,6 +1029,9 @@ class API:
        stream: bool,
        partial_images: int,
        bench: bool,
+        quality: Literal["high", "medium", "low"],
+        output_format: Literal["png", "jpeg", "webp"],
+        advanced_params: AdvancedImageParams | None,
    ) -> ImageEdits:
        """Prepare and send an image edits command with chunked image upload."""
        resolved_model = await self._validate_image_model(model)
@@ -1052,6 +1060,9 @@ class API:
                stream=stream,
                partial_images=partial_images,
                bench=bench,
+                quality=quality,
+                output_format=output_format,
+                advanced_params=advanced_params,
            ),
        )

@@ -1086,12 +1097,22 @@ class API:
        input_fidelity: Literal["low", "high"] = Form("low"),
        stream: str = Form("false"),
        partial_images: str = Form("0"),
+        quality: Literal["high", "medium", "low"] = Form("medium"),
+        output_format: Literal["png", "jpeg", "webp"] = Form("png"),
+        advanced_params: str | None = Form(None),
    ) -> ImageGenerationResponse | StreamingResponse:
        """Handle image editing requests (img2img)."""
        # Parse string form values to proper types
        stream_bool = stream.lower() in ("true", "1", "yes")
        partial_images_int = int(partial_images) if partial_images.isdigit() else 0

+        parsed_advanced_params: AdvancedImageParams | None = None
+        if advanced_params:
+            with contextlib.suppress(Exception):
+                parsed_advanced_params = AdvancedImageParams.model_validate_json(
+                    advanced_params
+                )
+
        command = await self._send_image_edits_command(
            image=image,
            prompt=prompt,
@@ -1103,6 +1124,9 @@ class API:
            stream=stream_bool,
            partial_images=partial_images_int,
            bench=False,
+            quality=quality,
+            output_format=output_format,
+            advanced_params=parsed_advanced_params,
        )

        if stream_bool and partial_images_int > 0:
@@ -1133,8 +1157,18 @@ class API:
        size: str = Form("1024x1024"),
        response_format: Literal["url", "b64_json"] = Form("b64_json"),
        input_fidelity: Literal["low", "high"] = Form("low"),
+        quality: Literal["high", "medium", "low"] = Form("medium"),
+        output_format: Literal["png", "jpeg", "webp"] = Form("png"),
+        advanced_params: str | None = Form(None),
    ) -> BenchImageGenerationResponse:
        """Handle benchmark image editing requests with generation stats."""
+        parsed_advanced_params: AdvancedImageParams | None = None
+        if advanced_params:
+            with contextlib.suppress(Exception):
+                parsed_advanced_params = AdvancedImageParams.model_validate_json(
+                    advanced_params
+                )
+
        command = await self._send_image_edits_command(
            image=image,
            prompt=prompt,
@@ -1146,6 +1180,9 @@ class API:
            stream=False,
            partial_images=0,
            bench=True,
+            quality=quality,
+            output_format=output_format,
+            advanced_params=parsed_advanced_params,
        )

        return await self._collect_image_generation_with_stats(
@@ -1196,6 +1233,7 @@ class API:
            tg.start_soon(self._apply_state)
            tg.start_soon(self._pause_on_new_election)
            tg.start_soon(self._cleanup_expired_images)
+            tg.start_soon(self._state_catchup)
            print_startup_banner(self.port)
            await serve(
                cast(ASGIFramework, self.app),
@@ -1206,6 +1244,22 @@ class API:
        self.command_sender.close()
        self.global_event_receiver.close()

+    async def _state_catchup(self):
+        with self.state_catchup_receiver as states:
+            async for state in states:
+                if (
+                    self.state.last_event_applied_idx == -1
+                    and state.last_event_applied_idx > self.state.last_event_applied_idx
+                ):
+                    logger.info(
+                        f"API catching up state to idx {state.last_event_applied_idx}"
+                    )
+                    self.event_buffer.store = {}
+                    self.event_buffer.next_idx_to_release = (
+                        state.last_event_applied_idx + 1
+                    )
+                    self.state = state
+
    async def _apply_state(self):
        with self.global_event_receiver as events:
            async for f_event in events:
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -68,6 +68,8 @@ class Master:
        # Send events to the forwarder to be indexed (usually from command processing)
        # Ideally these would be MasterForwarderEvents but type system says no :(
        global_event_sender: Sender[ForwarderEvent],
+        # not a fan but - send the entire state to a node so it can catchup without the whole event log.
+        state_catchup_sender: Sender[State],
    ):
        self.state = State()
        self._tg: TaskGroup = anyio.create_task_group()
@@ -77,6 +79,7 @@ class Master:
        self.command_receiver = command_receiver
        self.local_event_receiver = local_event_receiver
        self.global_event_sender = global_event_sender
+        self.state_catchup_sender = state_catchup_sender
        send, recv = channel[Event]()
        self.event_sender: Sender[Event] = send
        self._loopback_event_receiver: Receiver[Event] = recv
@@ -84,7 +87,6 @@ class Master:
            local_event_receiver.clone_sender()
        )
        self._multi_buffer = MultiSourceBuffer[NodeId, Event]()
-        # TODO: not have this
        self._event_log: list[Event] = []

    async def run(self):
@@ -291,11 +293,17 @@ class Master:
                                    command.finished_command_id
                                ]
                        case RequestEventLog():
-                            # We should just be able to send everything, since other buffers will ignore old messages
-                            for i in range(command.since_idx, len(self._event_log)):
-                                await self._send_event(
-                                    IndexedEvent(idx=i, event=self._event_log[i])
+                            if command.since_idx == 0:
+                                # This is an optimization, and should not be relied upon in theory.
+                                logger.info(
+                                    f"Master sending catchup state for index {self.state.last_event_applied_idx}"
                                )
+                                await self.state_catchup_sender.send(self.state)
+                            else:
+                                for i in range(command.since_idx, len(self._event_log)):
+                                    await self._send_event(
+                                        IndexedEvent(idx=i, event=self._event_log[i])
+                                    )
                    for event in generated_events:
                        await self.event_sender.send(event)
                except ValueError as e:
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -27,6 +27,7 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
    MemoryUsage,
 )
+from exo.shared.types.state import State
 from exo.shared.types.tasks import ChatCompletion as ChatCompletionTask
 from exo.shared.types.tasks import TaskStatus
 from exo.shared.types.worker.instances import (
@@ -47,6 +48,7 @@ async def test_master():
    ge_sender, global_event_receiver = channel[ForwarderEvent]()
    command_sender, co_receiver = channel[ForwarderCommand]()
    local_event_sender, le_receiver = channel[ForwarderEvent]()
+    st_s, _st_r = channel[State]()

    all_events: list[IndexedEvent] = []

@@ -67,6 +69,7 @@ async def test_master():
        global_event_sender=ge_sender,
        local_event_receiver=le_receiver,
        command_receiver=co_receiver,
+        state_catchup_sender=st_s,
    )
    logger.info("run the master")
    async with anyio.create_task_group() as tg:
--- a/src/exo/routing/topics.py
+++ b/src/exo/routing/topics.py
@@ -7,6 +7,7 @@ from exo.shared.types.commands import ForwarderCommand
 from exo.shared.types.events import (
    ForwarderEvent,
 )
+from exo.shared.types.state import State
 from exo.utils.pydantic_ext import CamelCaseModel


@@ -45,3 +46,4 @@ ELECTION_MESSAGES = TypedTopic(
 CONNECTION_MESSAGES = TypedTopic(
    "connection_messages", PublishPolicy.Never, ConnectionMessage
 )
+STATE_CATCHUP = TypedTopic("state_catchup", PublishPolicy.Always, State)
--- a/src/exo/shared/types/worker/runner_response.py
+++ b/src/exo/shared/types/worker/runner_response.py
@@ -30,6 +30,7 @@ class ImageGenerationResponse(BaseRunnerResponse):
    image_data: bytes
    format: Literal["png", "jpeg", "webp"] = "png"
    stats: ImageGenerationStats | None = None
+    image_index: int = 0

    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
        for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
@@ -44,6 +45,7 @@ class PartialImageResponse(BaseRunnerResponse):
    format: Literal["png", "jpeg", "webp"] = "png"
    partial_index: int
    total_partials: int
+    image_index: int = 0

    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
        for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
--- a/src/exo/worker/engines/image/generate.py
+++ b/src/exo/worker/engines/image/generate.py
@@ -75,19 +75,20 @@ def generate_image(
    intermediate images, then ImageGenerationResponse for the final image.

    Yields:
-        PartialImageResponse for intermediate images (if partial_images > 0)
-        ImageGenerationResponse for the final complete image
+        PartialImageResponse for intermediate images (if partial_images > 0, first image only)
+        ImageGenerationResponse for final complete images
    """
    width, height = parse_size(task.size)
    quality: Literal["low", "medium", "high"] = task.quality or "medium"

    advanced_params = task.advanced_params
    if advanced_params is not None and advanced_params.seed is not None:
-        seed = advanced_params.seed
+        base_seed = advanced_params.seed
    else:
-        seed = random.randint(0, 2**32 - 1)
+        base_seed = random.randint(0, 2**32 - 1)

    is_bench = getattr(task, "bench", False)
+    num_images = task.n or 1

    generation_start_time: float = 0.0

@@ -95,7 +96,11 @@ def generate_image(
        mx.reset_peak_memory()
        generation_start_time = time.perf_counter()

-    partial_images = task.partial_images or (3 if task.stream else 0)
+    partial_images = (
+        task.partial_images
+        if task.partial_images is not None
+        else (3 if task.stream else 0)
+    )

    image_path: Path | None = None

@@ -105,72 +110,81 @@ def generate_image(
            image_path = Path(tmpdir) / "input.png"
            image_path.write_bytes(base64.b64decode(task.image_data))

-        # Iterate over generator results
-        for result in model.generate(
-            prompt=task.prompt,
-            height=height,
-            width=width,
-            quality=quality,
-            seed=seed,
-            image_path=image_path,
-            partial_images=partial_images,
-            advanced_params=advanced_params,
-        ):
-            if isinstance(result, tuple):
-                # Partial image: (Image, partial_index, total_partials)
-                image, partial_idx, total_partials = result
-                buffer = io.BytesIO()
-                image_format = task.output_format.upper()
-                if image_format == "JPG":
-                    image_format = "JPEG"
-                if image_format == "JPEG" and image.mode == "RGBA":
-                    image = image.convert("RGB")
-                image.save(buffer, format=image_format)
+        for image_num in range(num_images):
+            # Increment seed for each image to ensure unique results
+            current_seed = base_seed + image_num

-                yield PartialImageResponse(
-                    image_data=buffer.getvalue(),
-                    format=task.output_format,
-                    partial_index=partial_idx,
-                    total_partials=total_partials,
-                )
-            else:
-                image = result
+            for result in model.generate(
+                prompt=task.prompt,
+                height=height,
+                width=width,
+                quality=quality,
+                seed=current_seed,
+                image_path=image_path,
+                partial_images=partial_images,
+                advanced_params=advanced_params,
+            ):
+                if isinstance(result, tuple):
+                    # Partial image: (Image, partial_index, total_partials)
+                    image, partial_idx, total_partials = result
+                    buffer = io.BytesIO()
+                    image_format = task.output_format.upper()
+                    if image_format == "JPG":
+                        image_format = "JPEG"
+                    if image_format == "JPEG" and image.mode == "RGBA":
+                        image = image.convert("RGB")
+                    image.save(buffer, format=image_format)

-                stats: ImageGenerationStats | None = None
-                if is_bench:
-                    generation_end_time = time.perf_counter()
-                    total_generation_time = generation_end_time - generation_start_time
-
-                    num_inference_steps = model.get_steps_for_quality(quality)
-
-                    seconds_per_step = (
-                        total_generation_time / num_inference_steps
-                        if num_inference_steps > 0
-                        else 0.0
+                    yield PartialImageResponse(
+                        image_data=buffer.getvalue(),
+                        format=task.output_format,
+                        partial_index=partial_idx,
+                        total_partials=total_partials,
+                        image_index=image_num,
                    )
+                else:
+                    image = result

-                    peak_memory_gb = mx.get_peak_memory() / (1024**3)
+                    # Only include stats on the final image
+                    stats: ImageGenerationStats | None = None
+                    if is_bench and image_num == num_images - 1:
+                        generation_end_time = time.perf_counter()
+                        total_generation_time = (
+                            generation_end_time - generation_start_time
+                        )

-                    stats = ImageGenerationStats(
-                        seconds_per_step=seconds_per_step,
-                        total_generation_time=total_generation_time,
-                        num_inference_steps=num_inference_steps,
-                        num_images=task.n or 1,
-                        image_width=width,
-                        image_height=height,
-                        peak_memory_usage=Memory.from_gb(peak_memory_gb),
+                        num_inference_steps = model.get_steps_for_quality(quality)
+                        total_steps = num_inference_steps * num_images
+
+                        seconds_per_step = (
+                            total_generation_time / total_steps
+                            if total_steps > 0
+                            else 0.0
+                        )
+
+                        peak_memory_gb = mx.get_peak_memory() / (1024**3)
+
+                        stats = ImageGenerationStats(
+                            seconds_per_step=seconds_per_step,
+                            total_generation_time=total_generation_time,
+                            num_inference_steps=num_inference_steps,
+                            num_images=num_images,
+                            image_width=width,
+                            image_height=height,
+                            peak_memory_usage=Memory.from_gb(peak_memory_gb),
+                        )
+
+                    buffer = io.BytesIO()
+                    image_format = task.output_format.upper()
+                    if image_format == "JPG":
+                        image_format = "JPEG"
+                    if image_format == "JPEG" and image.mode == "RGBA":
+                        image = image.convert("RGB")
+                    image.save(buffer, format=image_format)
+
+                    yield ImageGenerationResponse(
+                        image_data=buffer.getvalue(),
+                        format=task.output_format,
+                        stats=stats,
+                        image_index=image_num,
                    )
-
-                buffer = io.BytesIO()
-                image_format = task.output_format.upper()
-                if image_format == "JPG":
-                    image_format = "JPEG"
-                if image_format == "JPEG" and image.mode == "RGBA":
-                    image = image.convert("RGB")
-                image.save(buffer, format=image_format)
-
-                yield ImageGenerationResponse(
-                    image_data=buffer.getvalue(),
-                    format=task.output_format,
-                    stats=stats,
-                )
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -67,9 +67,8 @@ class Worker:
        connection_message_receiver: Receiver[ConnectionMessage],
        global_event_receiver: Receiver[ForwarderEvent],
        local_event_sender: Sender[ForwarderEvent],
-        # This is for requesting updates. It doesn't need to be a general command sender right now,
-        # but I think it's the correct way to be thinking about commands
        command_sender: Sender[ForwarderCommand],
+        state_catchup_receiver: Receiver[State],
    ):
        self.node_id: NodeId = node_id
        self.session_id: SessionId = session_id
@@ -79,6 +78,7 @@ class Worker:

        self.global_event_receiver = global_event_receiver
        self.local_event_sender = local_event_sender
+        self.state_catchup_receiver = state_catchup_receiver
        self.local_event_index = 0
        self.command_sender = command_sender
        self.connection_message_receiver = connection_message_receiver
@@ -117,6 +117,7 @@ class Worker:
            tg.start_soon(self._event_applier)
            tg.start_soon(self._forward_events)
            tg.start_soon(self._poll_connection_updates)
+            tg.start_soon(self._check_catchup_state)

        # Actual shutdown code - waits for all tasks to complete before executing.
        self.local_event_sender.close()
@@ -135,6 +136,22 @@ class Worker:
                    )
                )

+    async def _check_catchup_state(self):
+        with self.state_catchup_receiver as states:
+            async for state in states:
+                if (
+                    self.state.last_event_applied_idx == -1
+                    and state.last_event_applied_idx > self.state.last_event_applied_idx
+                ):
+                    logger.info(
+                        f"Worker catching up state to idx {state.last_event_applied_idx}"
+                    )
+                    self.event_buffer.store = {}
+                    self.event_buffer.next_idx_to_release = (
+                        state.last_event_applied_idx + 1
+                    )
+                    self.state = state
+
    async def _event_applier(self):
        with self.global_event_receiver as events:
            async for f_event in events:
@@ -342,10 +359,7 @@ class Worker:
        # We request all events after (and including) the missing index.
        # This function is started whenever we receive an event that is out of sequence.
        # It is cancelled as soon as we receiver an event that is in sequence.
-
-        if since_idx < 0:
-            logger.warning(f"Negative value encountered for nack request {since_idx=}")
-            since_idx = 0
+        assert since_idx >= 0

        with CancelScope() as scope:
            self._nack_cancel_scope = scope
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -612,7 +612,7 @@ def _process_image_response(
        command_id=command_id,
        model_id=shard_metadata.model_card.model_id,
        event_sender=event_sender,
-        image_index=response.partial_index if is_partial else image_index,
+        image_index=response.image_index,
        is_partial=is_partial,
        partial_index=response.partial_index if is_partial else None,
        total_partials=response.total_partials if is_partial else None,