Eval noise before CFG exchange

## Motivation

The dashboard UI attempted to parse all image generation responses as
SSE streams, even when streaming was disabled. This broke non-streaming
image generation.

## Changes

- Parse JSON responses directly when not streaming, use SSE parser only
when stream=true AND partialImages > 0
- explicitly disable partial images when not streaming

## Why It Works

Both API and dashboard now use the same condition (stream &&
partialImages > 0) to determine response format, ensuring correct
parsing.

## Test Plan

### Manual Testing

Non-streamed image generation results appear in the UI. Streamed image
generation still works

.github/actions/typecheck/action.yml

@@ -1,12 +0,0 @@
-name: Type Check
-
-description: "Run type checker"
-
-runs:
-  using: "composite"
-  steps:
-    - name: Run type checker
-      run: |
-        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just sync
-        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just check
-      shell: bash

.github/workflows/pipeline.yml

@@ -26,73 +26,14 @@ jobs:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
 
-      - name: Configure git user
-        run: |
-          git config --local user.email "github-actions@users.noreply.github.com"
-          git config --local user.name  "github-actions bot"
-        shell: bash
+      - name: Load nix develop environment
+        run: nix run github:nicknovitski/nix-develop/v1
 
-      - name: Pull LFS files
-        run: |
-          echo "Pulling Git LFS files..."
-          git lfs pull
-        shell: bash
+      - name: Sync dependencies
+        run: uv sync --all-packages
 
-      - name: Setup Nix Environment
-        run: |
-          echo "Checking for nix installation..."
-          
-          # Check if nix binary exists directly
-          if [ -f /nix/var/nix/profiles/default/bin/nix ]; then
-            echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix"
-            export PATH="/nix/var/nix/profiles/default/bin:$PATH"
-            echo "PATH=$PATH" >> $GITHUB_ENV
-            nix --version
-          elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
-            echo "Found nix profile script, sourcing..."
-            source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
-            nix --version
-          elif command -v nix >/dev/null 2>&1; then
-            echo "Nix already in PATH"
-            nix --version
-          else
-            echo "Nix not found. Debugging info:"
-            echo "Contents of /nix/var/nix/profiles/default/:"
-            ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found"
-            echo "Contents of /nix/var/nix/profiles/default/bin/:"
-            ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found"
-            exit 1
-          fi
-        shell: bash
-
-      - name: Configure basedpyright include for local MLX
-        run: |
-          RUNNER_LABELS='${{ toJSON(runner.labels) }}'
-          if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then
-            if [ -d "/Users/Shared/mlx" ]; then
-              echo "Updating [tool.basedpyright].include to use /Users/Shared/mlx"
-              awk '
-                BEGIN { in=0 }
-                /^\[tool\.basedpyright\]/ { in=1; print; next }
-                in && /^\[/ { in=0 }  # next section
-                in && /^[ \t]*include[ \t]*=/ {
-                  print "include = [\"/Users/Shared/mlx\"]"
-                  next
-                }
-                { print }
-              ' pyproject.toml > pyproject.toml.tmp && mv pyproject.toml.tmp pyproject.toml
-
-              echo "New [tool.basedpyright] section:"
-              sed -n '/^\[tool\.basedpyright\]/,/^\[/p' pyproject.toml | sed '$d' || true
-            else
-              echo "local_mlx tag present but /Users/Shared/mlx not found; leaving pyproject unchanged."
-            fi
-          else
-            echo "Runner does not have 'local_mlx' tag; leaving pyproject unchanged."
-          fi
-        shell: bash
-
-      - uses: ./.github/actions/typecheck
+      - name: Run type checker
+        run: uv run basedpyright --project pyproject.toml
 
   nix:
     name: Build and check (${{ matrix.system }})
@@ -191,3 +132,14 @@ jobs:
 
       - name: Run nix flake check
         run: nix flake check
+
+      - name: Run pytest (macOS only)
+        if: runner.os == 'macOS'
+        run: |
+          # Build the test environment (requires relaxed sandbox for uv2nix on macOS)
+          TEST_ENV=$(nix build '.#exo-test-env' --option sandbox relaxed --print-out-paths)
+
+          # Run pytest outside sandbox (needs GPU access for MLX)
+          export HOME="$RUNNER_TEMP"
+          export EXO_TESTS=1
+          $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib

app/EXO/EXO.xcodeproj/project.pbxproj

@@ -342,6 +342,8 @@
 				SDKROOT = macosx;
 				SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)";
 				SWIFT_OPTIMIZATION_LEVEL = "-Onone";
+				SWIFT_TREAT_WARNINGS_AS_ERRORS = YES;
+				GCC_TREAT_WARNINGS_AS_ERRORS = YES;
 			};
 			name = Debug;
 		};
@@ -397,6 +399,8 @@
 				MTL_FAST_MATH = YES;
 				SDKROOT = macosx;
 				SWIFT_COMPILATION_MODE = wholemodule;
+				SWIFT_TREAT_WARNINGS_AS_ERRORS = YES;
+				GCC_TREAT_WARNINGS_AS_ERRORS = YES;
 			};
 			name = Release;
 		};

dashboard/package-lock.json

@@ -865,7 +865,6 @@
 			"integrity": "sha512-oH8tXw7EZnie8FdOWYrF7Yn4IKrqTFHhXvl8YxXxbKwTMcD/5NNCryUSEXRk2ZR4ojnub0P8rNrsVGHXWqIDtA==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"@standard-schema/spec": "^1.0.0",
 				"@sveltejs/acorn-typescript": "^1.0.5",
@@ -905,7 +904,6 @@
 			"integrity": "sha512-Y1Cs7hhTc+a5E9Va/xwKlAJoariQyHY+5zBgCZg4PFWNYQ1nMN9sjK1zhw1gK69DuqVP++sht/1GZg1aRwmAXQ==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"@sveltejs/vite-plugin-svelte-inspector": "^4.0.1",
 				"debug": "^4.4.1",
@@ -1522,7 +1520,6 @@
 			"integrity": "sha512-LCCV0HdSZZZb34qifBsyWlUmok6W7ouER+oQIGBScS8EsZsQbrtFTUrDX4hOl+CS6p7cnNC4td+qrSVGSCTUfQ==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"undici-types": "~6.21.0"
 			}
@@ -1532,7 +1529,6 @@
 			"resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz",
 			"integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
 			"license": "MIT",
-			"peer": true,
 			"bin": {
 				"acorn": "bin/acorn"
 			},
@@ -1945,7 +1941,6 @@
 			"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
 			"dev": true,
 			"license": "ISC",
-			"peer": true,
 			"engines": {
 				"node": ">=12"
 			}
@@ -2653,7 +2648,6 @@
 			"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"engines": {
 				"node": ">=12"
 			},
@@ -2696,7 +2690,6 @@
 			"integrity": "sha512-UOnG6LftzbdaHZcKoPFtOcCKztrQ57WkHDeRD9t/PTQtmT0NHSeWWepj6pS0z/N7+08BHFDQVUrfmfMRcZwbMg==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"bin": {
 				"prettier": "bin/prettier.cjs"
 			},
@@ -2869,7 +2862,6 @@
 			"resolved": "https://registry.npmjs.org/svelte/-/svelte-5.45.3.tgz",
 			"integrity": "sha512-ngKXNhNvwPzF43QqEhDOue7TQTrG09em1sd4HBxVF0Wr2gopAmdEWan+rgbdgK4fhBtSOTJO8bYU4chUG7VXZQ==",
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"@jridgewell/remapping": "^2.3.4",
 				"@jridgewell/sourcemap-codec": "^1.5.0",
@@ -3014,7 +3006,6 @@
 			"integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==",
 			"dev": true,
 			"license": "Apache-2.0",
-			"peer": true,
 			"bin": {
 				"tsc": "bin/tsc",
 				"tsserver": "bin/tsserver"
@@ -3036,7 +3027,6 @@
 			"integrity": "sha512-+Oxm7q9hDoLMyJOYfUYBuHQo+dkAloi33apOPP56pzj+vsdJDzr+j1NISE5pyaAuKL4A3UD34qd0lx5+kfKp2g==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"esbuild": "^0.25.0",
 				"fdir": "^6.4.4",

dashboard/src/lib/stores/app.svelte.ts

@@ -173,6 +173,11 @@ export interface PlacementPreviewResponse {
   previews: PlacementPreview[];
 }
 
+interface ImageApiResponse {
+  created: number;
+  data: Array<{ b64_json?: string; url?: string }>;
+}
+
 interface RawStateResponse {
   topology?: RawTopology;
   instances?: Record<
@@ -2095,107 +2100,137 @@ class AppStore {
         throw new Error(`API error: ${response.status} - ${errorText}`);
       }
 
-      const reader = response.body?.getReader();
-      if (!reader) {
-        throw new Error("No response body");
-      }
+      // Streaming requires both stream=true AND partialImages > 0
+      const isStreaming = params.stream && params.partialImages > 0;
 
-      interface ImageGenerationChunk {
-        data?: { b64_json?: string };
-        format?: string;
-        type?: "partial" | "final";
-        image_index?: number;
-        partial_index?: number;
-        total_partials?: number;
-      }
+      if (!isStreaming) {
+        // Non-streaming: parse JSON response directly
+        const jsonResponse = (await response.json()) as ImageApiResponse;
+        const format = params.outputFormat || "png";
+        const mimeType = `image/${format}`;
 
-      const numImages = params.numImages;
+        const attachments: MessageAttachment[] = jsonResponse.data
+          .filter((img) => img.b64_json)
+          .map((img, index) => ({
+            type: "generated-image" as const,
+            name: `generated-image-${index + 1}.${format}`,
+            preview: `data:${mimeType};base64,${img.b64_json}`,
+            mimeType,
+          }));
 
-      await this.parseSSEStream<ImageGenerationChunk>(
-        reader,
-        targetConversationId,
-        (parsed) => {
-          const imageData = parsed.data?.b64_json;
+        this.updateConversationMessage(
+          targetConversationId,
+          assistantMessage.id,
+          (msg) => {
+            msg.content = "";
+            msg.attachments = attachments;
+          },
+        );
+        this.syncActiveMessagesIfNeeded(targetConversationId);
+      } else {
+        // Streaming mode: use SSE parser
+        const reader = response.body?.getReader();
+        if (!reader) {
+          throw new Error("No response body");
+        }
 
-          if (imageData) {
-            const format = parsed.format || "png";
-            const mimeType = `image/${format}`;
-            const imageIndex = parsed.image_index ?? 0;
+        interface ImageGenerationChunk {
+          data?: { b64_json?: string };
+          format?: string;
+          type?: "partial" | "final";
+          image_index?: number;
+          partial_index?: number;
+          total_partials?: number;
+        }
 
-            if (parsed.type === "partial") {
-              // Update with partial image and progress
-              const partialNum = (parsed.partial_index ?? 0) + 1;
-              const totalPartials = parsed.total_partials ?? 3;
-              const progressText =
-                numImages > 1
-                  ? `Generating image ${imageIndex + 1}/${numImages}... ${partialNum}/${totalPartials}`
-                  : `Generating... ${partialNum}/${totalPartials}`;
+        const numImages = params.numImages;
 
-              const partialAttachment: MessageAttachment = {
-                type: "generated-image",
-                name: `generated-image.${format}`,
-                preview: `data:${mimeType};base64,${imageData}`,
-                mimeType,
-              };
+        await this.parseSSEStream<ImageGenerationChunk>(
+          reader,
+          targetConversationId,
+          (parsed) => {
+            const imageData = parsed.data?.b64_json;
 
-              this.updateConversationMessage(
-                targetConversationId,
-                assistantMessage.id,
-                (msg) => {
-                  msg.content = progressText;
-                  if (imageIndex === 0) {
-                    // First image - safe to replace attachments with partial preview
-                    msg.attachments = [partialAttachment];
-                  } else {
-                    // Subsequent images - keep existing finals, show partial at current position
-                    const existingAttachments = msg.attachments || [];
-                    // Keep only the completed final images (up to current imageIndex)
-                    const finals = existingAttachments.slice(0, imageIndex);
-                    msg.attachments = [...finals, partialAttachment];
-                  }
-                },
-              );
-            } else if (parsed.type === "final") {
-              // 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 (imageData) {
+              const format = parsed.format || "png";
+              const mimeType = `image/${format}`;
+              const imageIndex = parsed.image_index ?? 0;
 
-              this.updateConversationMessage(
-                targetConversationId,
-                assistantMessage.id,
-                (msg) => {
-                  if (imageIndex === 0) {
-                    // First final image - replace any partial preview
-                    msg.attachments = [newAttachment];
-                  } else {
-                    // Subsequent images - keep previous finals, replace partial at current position
-                    const existingAttachments = msg.attachments || [];
-                    // Slice keeps indices 0 to imageIndex-1 (the previous final images)
-                    const previousFinals = existingAttachments.slice(
-                      0,
-                      imageIndex,
-                    );
-                    msg.attachments = [...previousFinals, newAttachment];
-                  }
+              if (parsed.type === "partial") {
+                // Update with partial image and progress
+                const partialNum = (parsed.partial_index ?? 0) + 1;
+                const totalPartials = parsed.total_partials ?? 3;
+                const progressText =
+                  numImages > 1
+                    ? `Generating image ${imageIndex + 1}/${numImages}... ${partialNum}/${totalPartials}`
+                    : `Generating... ${partialNum}/${totalPartials}`;
 
-                  // Update progress message for multiple images
-                  if (numImages > 1 && imageIndex < numImages - 1) {
-                    msg.content = `Generating image ${imageIndex + 2}/${numImages}...`;
-                  } else {
-                    msg.content = "";
-                  }
-                },
-              );
+                const partialAttachment: MessageAttachment = {
+                  type: "generated-image",
+                  name: `generated-image.${format}`,
+                  preview: `data:${mimeType};base64,${imageData}`,
+                  mimeType,
+                };
+
+                this.updateConversationMessage(
+                  targetConversationId,
+                  assistantMessage.id,
+                  (msg) => {
+                    msg.content = progressText;
+                    if (imageIndex === 0) {
+                      // First image - safe to replace attachments with partial preview
+                      msg.attachments = [partialAttachment];
+                    } else {
+                      // Subsequent images - keep existing finals, show partial at current position
+                      const existingAttachments = msg.attachments || [];
+                      // Keep only the completed final images (up to current imageIndex)
+                      const finals = existingAttachments.slice(0, imageIndex);
+                      msg.attachments = [...finals, partialAttachment];
+                    }
+                  },
+                );
+              } else if (parsed.type === "final") {
+                // 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,
+                };
+
+                this.updateConversationMessage(
+                  targetConversationId,
+                  assistantMessage.id,
+                  (msg) => {
+                    if (imageIndex === 0) {
+                      // First final image - replace any partial preview
+                      msg.attachments = [newAttachment];
+                    } else {
+                      // Subsequent images - keep previous finals, replace partial at current position
+                      const existingAttachments = msg.attachments || [];
+                      // Slice keeps indices 0 to imageIndex-1 (the previous final images)
+                      const previousFinals = existingAttachments.slice(
+                        0,
+                        imageIndex,
+                      );
+                      msg.attachments = [...previousFinals, newAttachment];
+                    }
+
+                    // Update progress message for multiple images
+                    if (numImages > 1 && imageIndex < numImages - 1) {
+                      msg.content = `Generating image ${imageIndex + 2}/${numImages}...`;
+                    } else {
+                      msg.content = "";
+                    }
+                  },
+                );
+              }
+
+              this.syncActiveMessagesIfNeeded(targetConversationId);
             }
-
-            this.syncActiveMessagesIfNeeded(targetConversationId);
-          }
-        },
-      );
+          },
+        );
+      }
     } catch (error) {
       console.error("Error generating image:", error);
       this.handleStreamingError(
@@ -2343,69 +2378,98 @@ class AppStore {
         throw new Error(`API error: ${apiResponse.status} - ${errorText}`);
       }
 
-      const reader = apiResponse.body?.getReader();
-      if (!reader) {
-        throw new Error("No response body");
-      }
+      // Streaming requires both stream=true AND partialImages > 0
+      const isStreaming = params.stream && params.partialImages > 0;
 
-      interface ImageEditChunk {
-        data?: { b64_json?: string };
-        format?: string;
-        type?: "partial" | "final";
-        partial_index?: number;
-        total_partials?: number;
-      }
+      if (!isStreaming) {
+        // Non-streaming: parse JSON response directly
+        const jsonResponse = (await apiResponse.json()) as ImageApiResponse;
+        const format = params.outputFormat || "png";
+        const mimeType = `image/${format}`;
+        const attachments: MessageAttachment[] = jsonResponse.data
+          .filter((img) => img.b64_json)
+          .map((img) => ({
+            type: "generated-image" as const,
+            name: `edited-image.${format}`,
+            preview: `data:${mimeType};base64,${img.b64_json}`,
+            mimeType,
+          }));
 
-      await this.parseSSEStream<ImageEditChunk>(
-        reader,
-        targetConversationId,
-        (parsed) => {
-          const imageData = parsed.data?.b64_json;
+        this.updateConversationMessage(
+          targetConversationId,
+          assistantMessage.id,
+          (msg) => {
+            msg.content = "";
+            msg.attachments = attachments;
+          },
+        );
+        this.syncActiveMessagesIfNeeded(targetConversationId);
+      } else {
+        // Streaming mode: use SSE parser
+        const reader = apiResponse.body?.getReader();
+        if (!reader) {
+          throw new Error("No response body");
+        }
 
-          if (imageData) {
-            const format = parsed.format || "png";
-            const mimeType = `image/${format}`;
-            if (parsed.type === "partial") {
-              // Update with partial image and progress
-              const partialNum = (parsed.partial_index ?? 0) + 1;
-              const totalPartials = parsed.total_partials ?? 3;
-              this.updateConversationMessage(
-                targetConversationId,
-                assistantMessage.id,
-                (msg) => {
-                  msg.content = `Editing... ${partialNum}/${totalPartials}`;
-                  msg.attachments = [
-                    {
-                      type: "generated-image",
-                      name: `edited-image.${format}`,
-                      preview: `data:${mimeType};base64,${imageData}`,
-                      mimeType,
-                    },
-                  ];
-                },
-              );
-            } else if (parsed.type === "final") {
-              // Final image
-              this.updateConversationMessage(
-                targetConversationId,
-                assistantMessage.id,
-                (msg) => {
-                  msg.content = "";
-                  msg.attachments = [
-                    {
-                      type: "generated-image",
-                      name: `edited-image.${format}`,
-                      preview: `data:${mimeType};base64,${imageData}`,
-                      mimeType,
-                    },
-                  ];
-                },
-              );
+        interface ImageEditChunk {
+          data?: { b64_json?: string };
+          format?: string;
+          type?: "partial" | "final";
+          partial_index?: number;
+          total_partials?: number;
+        }
+
+        await this.parseSSEStream<ImageEditChunk>(
+          reader,
+          targetConversationId,
+          (parsed) => {
+            const imageData = parsed.data?.b64_json;
+
+            if (imageData) {
+              const format = parsed.format || "png";
+              const mimeType = `image/${format}`;
+              if (parsed.type === "partial") {
+                // Update with partial image and progress
+                const partialNum = (parsed.partial_index ?? 0) + 1;
+                const totalPartials = parsed.total_partials ?? 3;
+                this.updateConversationMessage(
+                  targetConversationId,
+                  assistantMessage.id,
+                  (msg) => {
+                    msg.content = `Editing... ${partialNum}/${totalPartials}`;
+                    msg.attachments = [
+                      {
+                        type: "generated-image",
+                        name: `edited-image.${format}`,
+                        preview: `data:${mimeType};base64,${imageData}`,
+                        mimeType,
+                      },
+                    ];
+                  },
+                );
+              } else if (parsed.type === "final") {
+                // Final image
+                this.updateConversationMessage(
+                  targetConversationId,
+                  assistantMessage.id,
+                  (msg) => {
+                    msg.content = "";
+                    msg.attachments = [
+                      {
+                        type: "generated-image",
+                        name: `edited-image.${format}`,
+                        preview: `data:${mimeType};base64,${imageData}`,
+                        mimeType,
+                      },
+                    ];
+                  },
+                );
+              }
+              this.syncActiveMessagesIfNeeded(targetConversationId);
             }
-            this.syncActiveMessagesIfNeeded(targetConversationId);
-          }
-        },
-      );
+          },
+        );
+      }
     } catch (error) {
       console.error("Error editing image:", error);
       this.handleStreamingError(

flake.lock

@@ -21,7 +21,9 @@
           "nixpkgs"
         ],
         "purescript-overlay": "purescript-overlay",
-        "pyproject-nix": "pyproject-nix"
+        "pyproject-nix": [
+          "pyproject-nix"
+        ]
       },
       "locked": {
         "lastModified": 1765953015,
@@ -149,19 +151,44 @@
         "type": "github"
       }
     },
+    "pyproject-build-systems": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ],
+        "pyproject-nix": [
+          "pyproject-nix"
+        ],
+        "uv2nix": [
+          "uv2nix"
+        ]
+      },
+      "locked": {
+        "lastModified": 1763662255,
+        "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=",
+        "owner": "pyproject-nix",
+        "repo": "build-system-pkgs",
+        "rev": "042904167604c681a090c07eb6967b4dd4dae88c",
+        "type": "github"
+      },
+      "original": {
+        "owner": "pyproject-nix",
+        "repo": "build-system-pkgs",
+        "type": "github"
+      }
+    },
     "pyproject-nix": {
       "inputs": {
         "nixpkgs": [
-          "dream2nix",
           "nixpkgs"
         ]
       },
       "locked": {
-        "lastModified": 1763017646,
-        "narHash": "sha256-Z+R2lveIp6Skn1VPH3taQIuMhABg1IizJd8oVdmdHsQ=",
+        "lastModified": 1764134915,
+        "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
         "owner": "pyproject-nix",
         "repo": "pyproject.nix",
-        "rev": "47bd6f296502842643078d66128f7b5e5370790c",
+        "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
         "type": "github"
       },
       "original": {
@@ -178,7 +205,10 @@
         "flake-parts": "flake-parts",
         "nixpkgs": "nixpkgs",
         "nixpkgs-swift": "nixpkgs-swift",
-        "treefmt-nix": "treefmt-nix"
+        "pyproject-build-systems": "pyproject-build-systems",
+        "pyproject-nix": "pyproject-nix",
+        "treefmt-nix": "treefmt-nix",
+        "uv2nix": "uv2nix"
       }
     },
     "rust-analyzer-src": {
@@ -239,6 +269,29 @@
         "repo": "treefmt-nix",
         "type": "github"
       }
+    },
+    "uv2nix": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ],
+        "pyproject-nix": [
+          "pyproject-nix"
+        ]
+      },
+      "locked": {
+        "lastModified": 1767701098,
+        "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=",
+        "owner": "pyproject-nix",
+        "repo": "uv2nix",
+        "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da",
+        "type": "github"
+      },
+      "original": {
+        "owner": "pyproject-nix",
+        "repo": "uv2nix",
+        "type": "github"
+      }
     }
   },
   "root": "root",

flake.nix

@@ -24,6 +24,26 @@
     dream2nix = {
       url = "github:nix-community/dream2nix";
       inputs.nixpkgs.follows = "nixpkgs";
+      inputs.pyproject-nix.follows = "pyproject-nix";
+    };
+
+    # Python packaging with uv2nix
+    pyproject-nix = {
+      url = "github:pyproject-nix/pyproject.nix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+
+    uv2nix = {
+      url = "github:pyproject-nix/uv2nix";
+      inputs.pyproject-nix.follows = "pyproject-nix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+
+    pyproject-build-systems = {
+      url = "github:pyproject-nix/build-system-pkgs";
+      inputs.pyproject-nix.follows = "pyproject-nix";
+      inputs.uv2nix.follows = "uv2nix";
+      inputs.nixpkgs.follows = "nixpkgs";
     };
 
     # Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
@@ -48,6 +68,7 @@
         inputs.treefmt-nix.flakeModule
         ./dashboard/parts.nix
         ./rust/parts.nix
+        ./python/parts.nix
       ];
 
       perSystem =
@@ -88,12 +109,6 @@
             };
           };
 
-          checks.lint = pkgs.runCommand "lint-check" { } ''
-            export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
-            ${pkgs.ruff}/bin/ruff check ${inputs.self}/
-            touch $out
-          '';
-
           packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
             let
               uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);

python/parts.nix

@@ -0,0 +1,93 @@
+{ inputs, ... }:
+{
+  perSystem =
+    { config, self', pkgs, lib, system, ... }:
+    let
+      # Load workspace from uv.lock
+      workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
+        workspaceRoot = inputs.self;
+      };
+
+      # Create overlay from workspace
+      # Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build
+      overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
+
+      # Override overlay to inject Nix-built components
+      exoOverlay = final: prev: {
+        # Replace workspace exo_pyo3_bindings with Nix-built wheel
+        exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
+          pname = "exo-pyo3-bindings";
+          version = "0.1.0";
+          src = self'.packages.exo_pyo3_bindings;
+          # Install from pre-built wheel
+          nativeBuildInputs = [ final.pyprojectWheelHook ];
+          dontStrip = true;
+        };
+      };
+
+      python = pkgs.python313;
+
+      # Overlay to provide build systems and custom packages
+      buildSystemsOverlay = final: prev: {
+        # Use our pure Nix-built MLX with Metal support
+        mlx = self'.packages.mlx;
+
+        # mlx-lm is a git dependency that needs setuptools
+        mlx-lm = prev.mlx-lm.overrideAttrs (old: {
+          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+            final.setuptools
+          ];
+        });
+      };
+
+      pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
+        inherit python;
+      }).overrideScope (
+        lib.composeManyExtensions [
+          inputs.pyproject-build-systems.overlays.default
+          overlay
+          exoOverlay
+          buildSystemsOverlay
+        ]
+      );
+      exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default;
+
+      # Virtual environment with dev dependencies for testing
+      testVenv = pythonSet.mkVirtualEnv "exo-test-env" (
+        workspace.deps.default // {
+          exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
+        }
+      );
+
+      exoPackage = pkgs.runCommand "exo"
+        {
+          nativeBuildInputs = [ pkgs.makeWrapper ];
+        }
+        ''
+          mkdir -p $out/bin
+
+          # Create wrapper scripts
+          for script in exo exo-master exo-worker; do
+            makeWrapper ${exoVenv}/bin/$script $out/bin/$script \
+              --set DASHBOARD_DIR ${self'.packages.dashboard}
+          done
+        '';
+    in
+    {
+      # Python package only available on macOS (requires MLX/Metal)
+      packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
+        exo = exoPackage;
+        # Test environment for running pytest outside of Nix sandbox (needs GPU access)
+        exo-test-env = testVenv;
+      };
+
+      checks = {
+        # Ruff linting (works on all platforms)
+        lint = pkgs.runCommand "ruff-lint" { } ''
+          export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
+          ${pkgs.ruff}/bin/ruff check ${inputs.self}/
+          touch $out
+        '';
+      };
+    };
+}

src/exo/master/api.py

@@ -1,6 +1,7 @@
 import base64
 import contextlib
 import json
+import random
 import time
 from collections.abc import AsyncGenerator
 from http import HTTPStatus
@@ -112,6 +113,15 @@ def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None)
     return f"image/{image_format or 'png'}"
 
 
+def _ensure_seed(params: AdvancedImageParams | None) -> AdvancedImageParams:
+    """Ensure advanced params has a seed set for distributed consistency."""
+    if params is None:
+        return AdvancedImageParams(seed=random.randint(0, 2**32 - 1))
+    if params.seed is None:
+        return params.model_copy(update={"seed": random.randint(0, 2**32 - 1)})
+    return params
+
+
 def chunk_to_response(
     chunk: TokenChunk | ToolCallChunk, command_id: CommandId
 ) -> ChatCompletionResponse:
@@ -772,6 +782,9 @@ class API:
         with SSE-formatted events for partial and final images.
         """
         payload.model = await self._validate_image_model(payload.model)
+        payload = payload.model_copy(
+            update={"advanced_params": _ensure_seed(payload.advanced_params)}
+        )
 
         command = ImageGeneration(
             request_params=payload,
@@ -1020,6 +1033,9 @@ class API:
 
         payload.stream = False
         payload.partial_images = 0
+        payload = payload.model_copy(
+            update={"advanced_params": _ensure_seed(payload.advanced_params)}
+        )
 
         command = ImageGeneration(
             request_params=payload,
@@ -1051,6 +1067,7 @@ class API:
     ) -> ImageEdits:
         """Prepare and send an image edits command with chunked image upload."""
         resolved_model = await self._validate_image_model(model)
+        advanced_params = _ensure_seed(advanced_params)
 
         image_content = await image.read()
         image_data = base64.b64encode(image_content).decode("utf-8")

src/exo/master/placement_utils.py

@@ -94,20 +94,35 @@ def get_shard_assignments_for_pipeline_parallel(
     runner_to_shard: dict[RunnerId, ShardMetadata] = {}
     node_to_runner: dict[NodeId, RunnerId] = {}
 
+    # Determine CFG parallelism topology
+    # CFG parallel only for even node counts with CFG models (2+ nodes)
+    use_cfg_parallel = model_card.uses_cfg and world_size >= 2 and world_size % 2 == 0
+    cfg_world_size = 2 if use_cfg_parallel else 1
+    pipeline_world_size = world_size // cfg_world_size
+
+    # For CFG parallel, we only need to allocate layers for one pipeline group
+    # (both CFG groups run the same layers). Use the first pipeline group's nodes.
+    pipeline_node_ids = cycle.node_ids[:pipeline_world_size]
+    pipeline_memory = sum(
+        (node_memory[node_id].ram_available for node_id in pipeline_node_ids),
+        start=Memory(),
+    )
+
     layer_allocations = allocate_layers_proportionally(
         total_layers=total_layers,
         memory_fractions=[
-            node_memory[node_id].ram_available.in_bytes / cycle_memory.in_bytes
-            for node_id in cycle.node_ids
+            node_memory[node_id].ram_available.in_bytes / pipeline_memory.in_bytes
+            for node_id in pipeline_node_ids
         ],
     )
 
-    # Validate each node has sufficient memory for its assigned layers
-    memory_per_layer = model_card.storage_size.in_bytes / total_layers
-    for i, (node_id, node_layers) in enumerate(
-        zip(cycle.node_ids, layer_allocations, strict=True)
-    ):
-        required_memory = node_layers * memory_per_layer
+    # Validate each pipeline node has sufficient memory for its assigned layers
+    # Use integer arithmetic to avoid floating point precision issues
+    total_storage_bytes = model_card.storage_size.in_bytes
+    for i, node_id in enumerate(pipeline_node_ids):
+        node_layers = layer_allocations[i]
+        # Integer division then multiply to get conservative estimate
+        required_memory = (total_storage_bytes * node_layers) // total_layers
         available_memory = node_memory[node_id].ram_available.in_bytes
         if required_memory > available_memory:
             raise ValueError(
@@ -116,24 +131,69 @@ def get_shard_assignments_for_pipeline_parallel(
                 f"but only has {available_memory / (1024**3):.2f} GB available"
             )
 
-    layers_assigned = 0
-    for i, (node_id, node_layers) in enumerate(
-        zip(cycle.node_ids, layer_allocations, strict=True)
-    ):
+    # CFG group 0: pipeline ranks in ascending order (0, 1, 2, ...)
+    # CFG group 1: pipeline ranks in descending order (reversed)
+    # This places both "last stages" as ring neighbors for CFG exchange.
+    position_to_cfg_pipeline = [(0, r) for r in range(pipeline_world_size)] + [
+        (1, r) for r in reversed(range(pipeline_world_size))
+    ]
+
+    cfg_pipeline_to_device: dict[tuple[int, int], int] = {
+        (cfg_rank, pipeline_rank): i
+        for i, (cfg_rank, pipeline_rank) in enumerate(position_to_cfg_pipeline)
+    }
+
+    for i, node_id in enumerate(cycle.node_ids):
+        cfg_rank, pipeline_rank = position_to_cfg_pipeline[i]
+
+        layers_before = sum(layer_allocations[:pipeline_rank])
+        node_layers = layer_allocations[pipeline_rank]
+
+        is_first_stage = pipeline_rank == 0
+        is_last_stage = pipeline_rank == pipeline_world_size - 1
+
+        if is_last_stage:
+            next_pipeline_device = None
+        else:
+            next_pipeline_device = cfg_pipeline_to_device[(cfg_rank, pipeline_rank + 1)]
+
+        if is_first_stage:
+            prev_pipeline_device = None
+        else:
+            prev_pipeline_device = cfg_pipeline_to_device[(cfg_rank, pipeline_rank - 1)]
+
+        if is_last_stage and use_cfg_parallel:
+            other_cfg_rank = 1 - cfg_rank
+            cfg_peer_device = cfg_pipeline_to_device[(other_cfg_rank, pipeline_rank)]
+        else:
+            cfg_peer_device = None
+
+        first_pipeline_device = cfg_pipeline_to_device[(cfg_rank, 0)]
+        last_pipeline_device = cfg_pipeline_to_device[
+            (cfg_rank, pipeline_world_size - 1)
+        ]
+
         runner_id = RunnerId()
 
         shard = PipelineShardMetadata(
             model_card=model_card,
             device_rank=i,
             world_size=world_size,
-            start_layer=layers_assigned,
-            end_layer=layers_assigned + node_layers,
+            start_layer=layers_before,
+            end_layer=layers_before + node_layers,
             n_layers=total_layers,
+            cfg_rank=cfg_rank,
+            cfg_world_size=cfg_world_size,
+            explicit_pipeline_rank=pipeline_rank,
+            next_pipeline_device=next_pipeline_device,
+            prev_pipeline_device=prev_pipeline_device,
+            cfg_peer_device=cfg_peer_device,
+            first_pipeline_device=first_pipeline_device,
+            last_pipeline_device=last_pipeline_device,
         )
 
         runner_to_shard[runner_id] = shard
         node_to_runner[node_id] = runner_id
-        layers_assigned += node_layers
 
     shard_assignments = ShardAssignments(
         model_id=model_card.model_id,

src/exo/master/tests/test_placement_utils.py

@@ -5,6 +5,7 @@ from exo.master.placement_utils import (
     filter_cycles_by_memory,
     get_mlx_jaccl_coordinators,
     get_shard_assignments,
+    get_shard_assignments_for_pipeline_parallel,
     get_smallest_cycles,
 )
 from exo.master.tests.conftest import (
@@ -20,7 +21,7 @@ from exo.shared.types.profiling import (
     NodeNetworkInfo,
 )
 from exo.shared.types.topology import Connection, SocketConnection
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
 
 
 def test_filter_cycles_by_memory():
@@ -487,3 +488,195 @@ def test_get_shard_assignments_insufficient_memory_raises():
         get_shard_assignments(
             model_card, selected_cycle, Sharding.Pipeline, node_memory
         )
+
+
+class TestCfgParallelPlacement:
+    def _create_ring_topology(self, node_ids: list[NodeId]) -> Topology:
+        topology = Topology()
+        for node_id in node_ids:
+            topology.add_node(node_id)
+
+        for i, node_id in enumerate(node_ids):
+            next_node = node_ids[(i + 1) % len(node_ids)]
+            conn = Connection(
+                source=node_id,
+                sink=next_node,
+                edge=create_socket_connection(i + 1),
+            )
+            topology.add_connection(conn)
+
+        return topology
+
+    def test_two_nodes_cfg_model_uses_cfg_parallel(self):
+        """Two nodes with CFG model should use CFG parallel (no pipeline)."""
+        node_a = NodeId()
+        node_b = NodeId()
+
+        topology = self._create_ring_topology([node_a, node_b])
+        cycles = [c for c in topology.get_cycles() if len(c) == 2]
+        cycle = cycles[0]
+
+        node_memory = {
+            node_a: create_node_memory(1000 * 1024),
+            node_b: create_node_memory(1000 * 1024),
+        }
+
+        model_card = ModelCard(
+            model_id=ModelId("qwen-image-test"),
+            n_layers=60,
+            storage_size=Memory.from_kb(1000),
+            hidden_size=1,
+            supports_tensor=False,
+            uses_cfg=True,
+            tasks=[ModelTask.TextToImage],
+        )
+
+        assignments = get_shard_assignments_for_pipeline_parallel(
+            model_card, cycle, node_memory
+        )
+
+        shards = list(assignments.runner_to_shard.values())
+        assert len(shards) == 2
+
+        # Both nodes should have all layers (no pipeline split)
+        for shard in shards:
+            assert isinstance(shard, PipelineShardMetadata)
+            assert shard.start_layer == 0
+            assert shard.end_layer == 60
+            assert shard.cfg_world_size == 2
+
+        cfg_ranks = sorted(
+            s.cfg_rank for s in shards if isinstance(s, PipelineShardMetadata)
+        )
+        assert cfg_ranks == [0, 1]
+
+    def test_four_nodes_cfg_model_uses_hybrid(self):
+        """Four nodes with CFG model should use 2 CFG groups x 2 pipeline stages."""
+        nodes = [NodeId() for _ in range(4)]
+
+        topology = self._create_ring_topology(nodes)
+        cycles = [c for c in topology.get_cycles() if len(c) == 4]
+        cycle = cycles[0]
+
+        node_memory = {n: create_node_memory(1000 * 1024) for n in nodes}
+
+        model_card = ModelCard(
+            model_id=ModelId("qwen-image-test"),
+            n_layers=60,
+            storage_size=Memory.from_kb(1000),
+            hidden_size=1,
+            supports_tensor=False,
+            uses_cfg=True,
+            tasks=[ModelTask.TextToImage],
+        )
+
+        assignments = get_shard_assignments_for_pipeline_parallel(
+            model_card, cycle, node_memory
+        )
+
+        shards = list(assignments.runner_to_shard.values())
+        assert len(shards) == 4
+
+        for shard in shards:
+            assert isinstance(shard, PipelineShardMetadata)
+            assert shard.cfg_world_size == 2
+            assert shard.pipeline_world_size == 2
+
+        # Check we have 2 nodes in each CFG group
+        cfg_0_shards = [
+            s
+            for s in shards
+            if isinstance(s, PipelineShardMetadata) and s.cfg_rank == 0
+        ]
+        cfg_1_shards = [
+            s
+            for s in shards
+            if isinstance(s, PipelineShardMetadata) and s.cfg_rank == 1
+        ]
+        assert len(cfg_0_shards) == 2
+        assert len(cfg_1_shards) == 2
+
+        # Both CFG groups should have the same layer assignments
+        cfg_0_layers = [(s.start_layer, s.end_layer) for s in cfg_0_shards]
+        cfg_1_layers = [(s.start_layer, s.end_layer) for s in cfg_1_shards]
+        assert sorted(cfg_0_layers) == sorted(cfg_1_layers)
+
+    def test_three_nodes_cfg_model_uses_sequential_cfg(self):
+        """Three nodes (odd) with CFG model should use sequential CFG."""
+        nodes = [NodeId() for _ in range(3)]
+
+        topology = self._create_ring_topology(nodes)
+        cycles = [c for c in topology.get_cycles() if len(c) == 3]
+        cycle = cycles[0]
+
+        node_memory = {n: create_node_memory(1000 * 1024) for n in nodes}
+
+        model_card = ModelCard(
+            model_id=ModelId("qwen-image-test"),
+            n_layers=60,
+            storage_size=Memory.from_kb(1000),
+            hidden_size=1,
+            supports_tensor=False,
+            uses_cfg=True,
+            tasks=[ModelTask.TextToImage],
+        )
+
+        assignments = get_shard_assignments_for_pipeline_parallel(
+            model_card, cycle, node_memory
+        )
+
+        shards = list(assignments.runner_to_shard.values())
+        assert len(shards) == 3
+
+        for shard in shards:
+            assert isinstance(shard, PipelineShardMetadata)
+            # cfg_world_size = 1 means sequential CFG
+            assert shard.cfg_world_size == 1
+            assert shard.cfg_rank == 0
+
+    def test_two_nodes_non_cfg_model_uses_pipeline(self):
+        """Two nodes with non-CFG model should use pure pipeline."""
+        node_a = NodeId()
+        node_b = NodeId()
+
+        topology = self._create_ring_topology([node_a, node_b])
+        cycles = [c for c in topology.get_cycles() if len(c) == 2]
+        cycle = cycles[0]
+
+        node_memory = {
+            node_a: create_node_memory(1000 * 1024),
+            node_b: create_node_memory(1000 * 1024),
+        }
+
+        model_card = ModelCard(
+            model_id=ModelId("flux-test"),
+            n_layers=57,
+            storage_size=Memory.from_kb(1000),
+            hidden_size=1,
+            supports_tensor=False,
+            uses_cfg=False,  # Non-CFG model
+            tasks=[ModelTask.TextToImage],
+        )
+
+        assignments = get_shard_assignments_for_pipeline_parallel(
+            model_card, cycle, node_memory
+        )
+
+        shards = list(assignments.runner_to_shard.values())
+        assert len(shards) == 2
+
+        for shard in shards:
+            assert isinstance(shard, PipelineShardMetadata)
+            # cfg_world_size = 1 means no CFG parallel
+            assert shard.cfg_world_size == 1
+            assert shard.cfg_rank == 0
+
+        # Should have actual layer sharding (pipeline)
+        layer_ranges = sorted(
+            (s.start_layer, s.end_layer)
+            for s in shards
+            if isinstance(s, PipelineShardMetadata)
+        )
+        # First shard starts at 0, last shard ends at 57
+        assert layer_ranges[0][0] == 0
+        assert layer_ranges[-1][1] == 57

src/exo/shared/models/model_cards.py

@@ -47,6 +47,7 @@ class ModelCard(CamelCaseModel):
     supports_tensor: bool
     tasks: list[ModelTask]
     components: list[ComponentInfo] | None = None
+    uses_cfg: bool = False
 
     @field_validator("tasks", mode="before")
     @classmethod
@@ -562,6 +563,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
         hidden_size=1,
         supports_tensor=False,
         tasks=[ModelTask.TextToImage],
+        uses_cfg=True,
         components=[
             ComponentInfo(
                 component_name="text_encoder",
@@ -596,6 +598,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
         hidden_size=1,
         supports_tensor=False,
         tasks=[ModelTask.ImageToImage],
+        uses_cfg=True,
         components=[
             ComponentInfo(
                 component_name="text_encoder",
@@ -681,6 +684,7 @@ def _generate_image_model_quant_variants(
             hidden_size=base_card.hidden_size,
             supports_tensor=base_card.supports_tensor,
             tasks=base_card.tasks,
+            uses_cfg=base_card.uses_cfg,
             components=with_transformer_size(transformer_bytes),
         )
     }
@@ -700,6 +704,7 @@ def _generate_image_model_quant_variants(
             hidden_size=base_card.hidden_size,
             supports_tensor=base_card.supports_tensor,
             tasks=base_card.tasks,
+            uses_cfg=base_card.uses_cfg,
             components=with_transformer_size(quant_transformer_bytes),
         )
 

src/exo/shared/types/worker/shards.py

@@ -57,8 +57,62 @@ class PipelineShardMetadata(BaseShardMetadata):
 
     Layers are represented as a half-open interval [start_layer, end_layer),
     where start_layer is inclusive and end_layer is exclusive.
+
+    CFG parallelism fields:
+    - cfg_rank: 0 = positive branch, 1 = negative branch (or 0 if no CFG parallel)
+    - cfg_world_size: 1 = sequential CFG, 2 = parallel CFG
+
+    Communication rank fields (explicit to support ring topology):
+    - next_pipeline_device: device to send to in pipeline forward pass
+    - prev_pipeline_device: device to receive from in pipeline forward pass
+    - cfg_peer_device: device for CFG exchange (last stage only)
+    - first_pipeline_device: device of first stage in same CFG group (for latent return)
     """
 
+    cfg_rank: int = 0
+    cfg_world_size: int = 1
+
+    # Explicit pipeline position (CFG group 1 uses reversed pipeline order)
+    explicit_pipeline_rank: int | None = None
+
+    next_pipeline_device: int | None = None
+    prev_pipeline_device: int | None = None
+    cfg_peer_device: int | None = None
+    first_pipeline_device: int | None = None
+    last_pipeline_device: int | None = None
+
+    @property
+    def pipeline_world_size(self) -> int:
+        return self.world_size // self.cfg_world_size
+
+    @property
+    def pipeline_rank(self) -> int:
+        if self.explicit_pipeline_rank is not None:
+            return self.explicit_pipeline_rank
+        return self.device_rank % self.pipeline_world_size
+
+    @property
+    def is_pipeline_first(self) -> bool:
+        return self.pipeline_rank == 0
+
+    @property
+    def is_pipeline_last(self) -> bool:
+        return self.pipeline_rank == self.pipeline_world_size - 1
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                self.model_card.model_id,
+                self.start_layer,
+                self.end_layer,
+                self.n_layers,
+                self.device_rank,
+                self.world_size,
+                self.cfg_rank,
+                self.cfg_world_size,
+            )
+        )
+
 
 class TensorShardMetadata(BaseShardMetadata):
     pass

src/exo/worker/engines/image/distributed_model.py

@@ -37,7 +37,12 @@ class DistributedImageModel:
         config = get_config_for_model(model_id)
         adapter = create_adapter_for_model(config, model_id, local_path, quantize)
 
-        if group is not None:
+        has_layer_sharding = (
+            shard_metadata.start_layer != 0
+            or shard_metadata.end_layer != shard_metadata.n_layers
+        )
+
+        if group is not None and has_layer_sharding:
             adapter.slice_transformer_blocks(
                 start_layer=shard_metadata.start_layer,
                 end_layer=shard_metadata.end_layer,

src/exo/worker/engines/image/generate.py

@@ -98,8 +98,8 @@ def generate_image(
 
     partial_images = (
         task.partial_images
-        if task.partial_images is not None
-        else (3 if task.stream else 0)
+        if task.partial_images is not None and task.stream is not None and task.stream
+        else 0
     )
 
     image_path: Path | None = None

src/exo/worker/engines/image/models/base.py

@@ -86,6 +86,27 @@ class PromptData(ABC):
         """
         ...
 
+    @abstractmethod
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        """Get embeddings for a single CFG branch (positive or negative).
+
+        Used for sequential CFG and CFG parallel modes where we process
+        one branch at a time instead of batching.
+
+        Args:
+            positive: True for positive prompt, False for negative prompt
+
+        Returns:
+            Tuple of:
+            - embeds: [1, seq, hidden] prompt embeddings
+            - mask: [1, seq] attention mask or None
+            - pooled: [1, hidden] pooled embeddings or None
+            - conditioning_latents: [1, latent_seq, latent_dim] or None
+        """
+        ...
+
 
 class ModelAdapter(ABC, Generic[ModelT, TransformerT]):
     _config: ImageModelConfig

src/exo/worker/engines/image/models/flux/adapter.py

@@ -64,6 +64,12 @@ class FluxPromptData(PromptData):
     ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
         return None
 
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        """Flux doesn't use CFG, but we return positive data for compatibility."""
+        return (self._prompt_embeds, None, self._pooled_prompt_embeds, None)
+
 
 class FluxModelAdapter(ModelAdapter[Flux1, Transformer]):
     def __init__(

src/exo/worker/engines/image/models/qwen/adapter.py

@@ -133,6 +133,24 @@ class QwenPromptData(PromptData):
 
         return batched_embeds, batched_mask, None, cond_latents
 
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        if positive:
+            return (
+                self._prompt_embeds,
+                self._prompt_mask,
+                None,
+                self.conditioning_latents,
+            )
+        else:
+            return (
+                self._negative_prompt_embeds,
+                self._negative_prompt_mask,
+                None,
+                self.conditioning_latents,
+            )
+
 
 class QwenModelAdapter(ModelAdapter[QwenImage, QwenTransformer]):
     """Adapter for Qwen-Image model.

src/exo/worker/engines/image/models/qwen/edit_adapter.py

@@ -153,6 +153,24 @@ class QwenEditPromptData(PromptData):
 
         return batched_embeds, batched_mask, None, batched_cond_latents
 
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        if positive:
+            return (
+                self._prompt_embeds,
+                self._prompt_mask,
+                None,
+                self._conditioning_latents,
+            )
+        else:
+            return (
+                self._negative_prompt_embeds,
+                self._negative_prompt_mask,
+                None,
+                self._conditioning_latents,
+            )
+
 
 class QwenEditModelAdapter(ModelAdapter[QwenImageEdit, QwenTransformer]):
     """Adapter for Qwen-Image-Edit model.

src/exo/worker/engines/image/pipeline/runner.py

@@ -1,5 +1,7 @@
+from collections.abc import Iterator
+from dataclasses import dataclass
 from math import ceil
-from typing import Any, Optional
+from typing import Any, Optional, final
 
 import mlx.core as mx
 from mflux.models.common.config.config import Config
@@ -20,6 +22,16 @@ from exo.worker.engines.image.pipeline.block_wrapper import (
 )
 
 
+@final
+@dataclass
+class CfgBranch:
+    positive: bool
+    embeds: mx.array
+    mask: mx.array | None
+    pooled: mx.array | None
+    cond_latents: mx.array | None
+
+
 def calculate_patch_heights(
     latent_height: int, num_patches: int
 ) -> tuple[list[int], int]:
@@ -72,22 +84,11 @@ class DiffusionRunner:
         self.adapter = adapter
         self.group = group
 
-        if group is None:
-            self.rank = 0
-            self.world_size = 1
-            self.next_rank = 0
-            self.prev_rank = 0
-            self.start_layer = 0
-            self.end_layer = config.total_blocks
-        else:
-            self.rank = shard_metadata.device_rank
-            self.world_size = shard_metadata.world_size
-            self.next_rank = (self.rank + 1) % self.world_size
-            self.prev_rank = (self.rank - 1 + self.world_size) % self.world_size
-            self.start_layer = shard_metadata.start_layer
-            self.end_layer = shard_metadata.end_layer
+        self._init_cfg_topology(shard_metadata)
 
-        self.num_patches = num_patches if num_patches else max(1, self.world_size)
+        self.num_patches = (
+            num_patches if num_patches else max(1, self.pipeline_world_size)
+        )
 
         self.total_joint = config.joint_block_count
         self.total_single = config.single_block_count
@@ -97,6 +98,48 @@ class DiffusionRunner:
 
         self._compute_assigned_blocks()
 
+    def _init_cfg_topology(self, shard_metadata: PipelineShardMetadata) -> None:
+        """Initialize CFG and pipeline topology from shard metadata."""
+        if self.group is None:
+            self.rank = 0
+            self.world_size = 1
+            self.start_layer = 0
+            self.end_layer = self.config.total_blocks
+
+            self.cfg_rank = 0
+            self.cfg_world_size = 1
+            self.cfg_parallel = False
+
+            self.pipeline_world_size = 1
+            self.pipeline_rank = 0
+
+            self.next_pipeline_rank: int | None = None
+            self.prev_pipeline_rank: int | None = None
+            self.cfg_peer_rank: int | None = None
+            self.first_pipeline_rank: int = 0
+            self.last_pipeline_rank: int = 0
+        else:
+            self.rank = shard_metadata.device_rank
+            self.world_size = shard_metadata.world_size
+            self.start_layer = shard_metadata.start_layer
+            self.end_layer = shard_metadata.end_layer
+
+            self.cfg_rank = shard_metadata.cfg_rank
+            self.cfg_world_size = shard_metadata.cfg_world_size
+            self.cfg_parallel = self.cfg_world_size > 1
+
+            self.pipeline_world_size = shard_metadata.pipeline_world_size
+            self.pipeline_rank = shard_metadata.pipeline_rank
+
+            self.next_pipeline_rank = shard_metadata.next_pipeline_device
+            self.prev_pipeline_rank = shard_metadata.prev_pipeline_device
+            self.cfg_peer_rank = shard_metadata.cfg_peer_device
+
+            assert shard_metadata.first_pipeline_device is not None
+            assert shard_metadata.last_pipeline_device is not None
+            self.first_pipeline_rank = shard_metadata.first_pipeline_device
+            self.last_pipeline_rank = shard_metadata.last_pipeline_device
+
     def _compute_assigned_blocks(self) -> None:
         """Determine which joint/single blocks this stage owns."""
         start = self.start_layer
@@ -133,11 +176,11 @@ class DiffusionRunner:
 
     @property
     def is_first_stage(self) -> bool:
-        return self.rank == 0
+        return self.pipeline_rank == 0
 
     @property
     def is_last_stage(self) -> bool:
-        return self.rank == self.world_size - 1
+        return self.pipeline_rank == self.pipeline_world_size - 1
 
     @property
     def is_distributed(self) -> bool:
@@ -148,6 +191,97 @@ class DiffusionRunner:
             return self._guidance_override
         return self.config.guidance_scale
 
+    def _get_cfg_branches(self, prompt_data: PromptData) -> Iterator[CfgBranch]:
+        """Yield the CFG branches this node should process.
+
+        - No CFG: yields one branch (positive)
+        - CFG parallel: yields one branch (our assigned branch)
+        - Sequential CFG: yields two branches (positive, then negative)
+        """
+        if not self.adapter.needs_cfg:
+            embeds, mask, pooled, cond = prompt_data.get_cfg_branch_data(positive=True)
+            yield CfgBranch(
+                positive=True,
+                embeds=embeds,
+                mask=mask,
+                pooled=pooled,
+                cond_latents=cond,
+            )
+        elif self.cfg_parallel:
+            positive = self.cfg_rank == 0
+            embeds, mask, pooled, cond = prompt_data.get_cfg_branch_data(positive)
+            yield CfgBranch(
+                positive=positive,
+                embeds=embeds,
+                mask=mask,
+                pooled=pooled,
+                cond_latents=cond,
+            )
+        else:
+            pos_embeds, pos_mask, pos_pooled, pos_cond = (
+                prompt_data.get_cfg_branch_data(positive=True)
+            )
+            yield CfgBranch(
+                positive=True,
+                embeds=pos_embeds,
+                mask=pos_mask,
+                pooled=pos_pooled,
+                cond_latents=pos_cond,
+            )
+            neg_embeds, neg_mask, neg_pooled, neg_cond = (
+                prompt_data.get_cfg_branch_data(positive=False)
+            )
+            yield CfgBranch(
+                positive=False,
+                embeds=neg_embeds,
+                mask=neg_mask,
+                pooled=neg_pooled,
+                cond_latents=neg_cond,
+            )
+
+    def _combine_cfg_results(self, results: list[tuple[bool, mx.array]]) -> mx.array:
+        if len(results) == 1:
+            positive, noise = results[0]
+            if self.cfg_parallel and self.is_last_stage:
+                # TODO(ciaran): try to remove
+                mx.eval(noise)
+                return self._exchange_and_apply_guidance(noise, positive)
+            return noise
+
+        noise_neg = next(n for p, n in results if not p)
+        noise_pos = next(n for p, n in results if p)
+        return self._apply_guidance(noise_pos, noise_neg)
+
+    def _exchange_and_apply_guidance(
+        self, noise: mx.array, is_positive: bool
+    ) -> mx.array:
+        assert self.group is not None
+        assert self.cfg_peer_rank is not None
+
+        if is_positive:
+            noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
+            mx.async_eval(noise)
+            noise_neg = mx.distributed.recv_like(
+                noise, self.cfg_peer_rank, group=self.group
+            )
+            mx.eval(noise_neg)
+            noise_pos = noise
+        else:
+            noise_pos = mx.distributed.recv_like(
+                noise, self.cfg_peer_rank, group=self.group
+            )
+            mx.eval(noise_pos)
+            noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
+            mx.async_eval(noise)
+            noise_neg = noise
+
+        return self._apply_guidance(noise_pos, noise_neg)
+
+    def _apply_guidance(self, noise_pos: mx.array, noise_neg: mx.array) -> mx.array:
+        scale = self._get_effective_guidance_scale()
+        assert scale is not None
+        return self.adapter.apply_guidance(noise_pos, noise_neg, scale)
+
     def _ensure_wrappers(
         self,
         text_seq_len: int,
@@ -348,6 +482,7 @@ class DiffusionRunner:
                 ctx.in_loop(  # pyright: ignore[reportAny]
                     t=t,
                     latents=latents,
+                    time_steps=time_steps,
                 )
 
                 mx.eval(latents)
@@ -463,7 +598,9 @@ class DiffusionRunner:
     ) -> mx.array:
         if self.group is None:
             return self._single_node_step(t, config, latents, prompt_data)
-        elif t < config.init_time_step + num_sync_steps:
+        elif (
+            self.pipeline_world_size == 1 or t < config.init_time_step + num_sync_steps
+        ):
             return self._sync_pipeline_step(
                 t,
                 config,
@@ -487,42 +624,29 @@ class DiffusionRunner:
         prompt_data: PromptData,
     ) -> mx.array:
         cond_image_grid = prompt_data.cond_image_grid
-        needs_cfg = self.adapter.needs_cfg
+        results: list[tuple[bool, mx.array]] = []
+
+        for branch in self._get_cfg_branches(prompt_data):
+            # Reset caches before each branch to ensure no state contamination
+            self._reset_all_caches()
 
-        if needs_cfg:
-            batched_data = prompt_data.get_batched_cfg_data()
-            assert batched_data is not None, "CFG model must provide batched data"
-            prompt_embeds, encoder_mask, batched_pooled, cond_latents = batched_data
             pooled_embeds = (
-                batched_pooled if batched_pooled is not None else prompt_embeds
-            )
-            step_latents = mx.concatenate([latents, latents], axis=0)
-        else:
-            prompt_embeds = prompt_data.prompt_embeds
-            pooled_embeds = prompt_data.pooled_prompt_embeds
-            encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
-            cond_latents = prompt_data.conditioning_latents
-            step_latents = latents
-
-        noise = self._forward_pass(
-            step_latents,
-            prompt_embeds,
-            pooled_embeds,
-            t=t,
-            config=config,
-            encoder_hidden_states_mask=encoder_mask,
-            cond_image_grid=cond_image_grid,
-            conditioning_latents=cond_latents,
-        )
-
-        if needs_cfg:
-            noise_pos, noise_neg = mx.split(noise, 2, axis=0)
-            guidance_scale = self._get_effective_guidance_scale()
-            assert guidance_scale is not None
-            noise = self.adapter.apply_guidance(
-                noise_pos, noise_neg, guidance_scale=guidance_scale
+                branch.pooled if branch.pooled is not None else branch.embeds
             )
 
+            noise = self._forward_pass(
+                latents,
+                branch.embeds,
+                pooled_embeds,
+                t=t,
+                config=config,
+                encoder_hidden_states_mask=branch.mask,
+                cond_image_grid=cond_image_grid,
+                conditioning_latents=branch.cond_latents,
+            )
+            results.append((branch.positive, noise))
+
+        noise = self._combine_cfg_results(results)
         return config.scheduler.step(noise=noise, timestep=t, latents=latents)  # pyright: ignore[reportAny]
 
     def _create_patches(
@@ -573,7 +697,7 @@ class DiffusionRunner:
             )
 
         text_embeddings = self.adapter.compute_text_embeddings(
-            t, config, pooled_prompt_embeds
+            t, config, pooled_prompt_embeds, hidden_states=hidden_states
         )
         image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
             prompt_embeds,
@@ -585,16 +709,17 @@ class DiffusionRunner:
 
         if self.has_joint_blocks:
             if not self.is_first_stage:
+                assert self.prev_pipeline_rank is not None
                 hidden_states = mx.distributed.recv(
                     (batch_size, num_img_tokens, hidden_dim),
                     dtype,
-                    self.prev_rank,
+                    self.prev_pipeline_rank,
                     group=self.group,
                 )
                 encoder_hidden_states = mx.distributed.recv(
                     (batch_size, text_seq_len, hidden_dim),
                     dtype,
-                    self.prev_rank,
+                    self.prev_pipeline_rank,
                     group=self.group,
                 )
                 mx.eval(hidden_states, encoder_hidden_states)
@@ -619,27 +744,30 @@ class DiffusionRunner:
             if self.has_single_blocks or self.is_last_stage:
                 hidden_states = concatenated
             else:
+                assert self.next_pipeline_rank is not None
                 concatenated = mx.distributed.send(
-                    concatenated, self.next_rank, group=self.group
+                    concatenated, self.next_pipeline_rank, group=self.group
                 )
                 mx.async_eval(concatenated)
 
         elif self.has_joint_blocks and not self.is_last_stage:
             assert encoder_hidden_states is not None
+            assert self.next_pipeline_rank is not None
             hidden_states = mx.distributed.send(
-                hidden_states, self.next_rank, group=self.group
+                hidden_states, self.next_pipeline_rank, group=self.group
             )
             encoder_hidden_states = mx.distributed.send(
-                encoder_hidden_states, self.next_rank, group=self.group
+                encoder_hidden_states, self.next_pipeline_rank, group=self.group
             )
             mx.async_eval(hidden_states, encoder_hidden_states)
 
         if self.has_single_blocks:
             if not self.owns_concat_stage and not self.is_first_stage:
+                assert self.prev_pipeline_rank is not None
                 hidden_states = mx.distributed.recv(
                     (batch_size, text_seq_len + num_img_tokens, hidden_dim),
                     dtype,
-                    self.prev_rank,
+                    self.prev_pipeline_rank,
                     group=self.group,
                 )
                 mx.eval(hidden_states)
@@ -654,8 +782,9 @@ class DiffusionRunner:
                 )
 
             if not self.is_last_stage:
+                assert self.next_pipeline_rank is not None
                 hidden_states = mx.distributed.send(
-                    hidden_states, self.next_rank, group=self.group
+                    hidden_states, self.next_pipeline_rank, group=self.group
                 )
                 mx.async_eval(hidden_states)
 
@@ -678,75 +807,65 @@ class DiffusionRunner:
         kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
         prev_latents = hidden_states
-        needs_cfg = self.adapter.needs_cfg
         cond_image_grid = prompt_data.cond_image_grid
 
         scaled_hidden_states = config.scheduler.scale_model_input(hidden_states, t)  # pyright: ignore[reportAny]
         original_latent_tokens: int = scaled_hidden_states.shape[1]  # pyright: ignore[reportAny]
 
-        if needs_cfg:
-            batched_data = prompt_data.get_batched_cfg_data()
-            assert batched_data is not None, "CFG model must provide batched data"
-            prompt_embeds, encoder_mask, batched_pooled, cond_latents = batched_data
+        results: list[tuple[bool, mx.array]] = []
+
+        for branch in self._get_cfg_branches(prompt_data):
             pooled_embeds = (
-                batched_pooled if batched_pooled is not None else prompt_embeds
+                branch.pooled if branch.pooled is not None else branch.embeds
             )
-            step_latents = mx.concatenate(
-                [scaled_hidden_states, scaled_hidden_states], axis=0
+
+            cond_latents = branch.cond_latents
+            if cond_latents is not None:
+                num_img_tokens: int = original_latent_tokens + cond_latents.shape[1]
+            else:
+                num_img_tokens = original_latent_tokens
+
+            step_latents: mx.array = scaled_hidden_states  # pyright: ignore[reportAny]
+            if self.is_first_stage and cond_latents is not None:
+                step_latents = mx.concatenate([step_latents, cond_latents], axis=1)
+
+            text_seq_len = branch.embeds.shape[1]
+            self._ensure_wrappers(text_seq_len, branch.mask)
+
+            noise = self._run_sync_pass(
+                t,
+                config,
+                step_latents,
+                branch.embeds,
+                pooled_embeds,
+                branch.mask,
+                cond_image_grid,
+                kontext_image_ids,
+                num_img_tokens,
+                original_latent_tokens,
+                cond_latents,
             )
-        else:
-            prompt_embeds = prompt_data.prompt_embeds
-            pooled_embeds = prompt_data.pooled_prompt_embeds
-            encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
-            cond_latents = prompt_data.conditioning_latents
-            step_latents = scaled_hidden_states  # pyright: ignore[reportAny]
 
-        if cond_latents is not None:
-            num_img_tokens: int = original_latent_tokens + cond_latents.shape[1]
-        else:
-            num_img_tokens = original_latent_tokens
-
-        if self.is_first_stage and cond_latents is not None:
-            step_latents = mx.concatenate([step_latents, cond_latents], axis=1)
-
-        text_seq_len = prompt_embeds.shape[1]
-        self._ensure_wrappers(text_seq_len, encoder_mask)
-
-        noise = self._run_sync_pass(
-            t,
-            config,
-            step_latents,
-            prompt_embeds,
-            pooled_embeds,
-            encoder_mask,
-            cond_image_grid,
-            kontext_image_ids,
-            num_img_tokens,
-            original_latent_tokens,
-            cond_latents,
-        )
+            if self.is_last_stage:
+                assert noise is not None
+                results.append((branch.positive, noise))
 
         if self.is_last_stage:
-            assert noise is not None
-            if needs_cfg:
-                noise_pos, noise_neg = mx.split(noise, 2, axis=0)
-                guidance_scale = self._get_effective_guidance_scale()
-                assert guidance_scale is not None
-                noise = self.adapter.apply_guidance(
-                    noise_pos, noise_neg, guidance_scale
-                )
+            noise = self._combine_cfg_results(results)
 
             hidden_states = config.scheduler.step(  # pyright: ignore[reportAny]
                 noise=noise, timestep=t, latents=prev_latents
             )
 
             if not self.is_first_stage:
-                hidden_states = mx.distributed.send(hidden_states, 0, group=self.group)
+                hidden_states = mx.distributed.send(
+                    hidden_states, self.first_pipeline_rank, group=self.group
+                )
                 mx.async_eval(hidden_states)
 
         elif self.is_first_stage:
             hidden_states = mx.distributed.recv_like(
-                prev_latents, src=self.world_size - 1, group=self.group
+                prev_latents, src=self.last_pipeline_rank, group=self.group
             )
             mx.eval(hidden_states)
 
@@ -765,39 +884,10 @@ class DiffusionRunner:
         kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
         patch_latents, token_indices = self._create_patches(latents, config)
-        needs_cfg = self.adapter.needs_cfg
         cond_image_grid = prompt_data.cond_image_grid
 
-        if needs_cfg:
-            batched_data = prompt_data.get_batched_cfg_data()
-            assert batched_data is not None, "CFG model must provide batched data"
-            prompt_embeds, encoder_mask, batched_pooled, _ = batched_data
-            pooled_embeds = (
-                batched_pooled if batched_pooled is not None else prompt_embeds
-            )
-        else:
-            prompt_embeds = prompt_data.prompt_embeds
-            pooled_embeds = prompt_data.pooled_prompt_embeds
-            encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
-
-        text_seq_len = prompt_embeds.shape[1]
-        self._ensure_wrappers(text_seq_len, encoder_mask)
-        self._set_text_seq_len(text_seq_len)
-
-        if self.joint_block_wrappers:
-            for wrapper in self.joint_block_wrappers:
-                wrapper.set_encoder_mask(encoder_mask)
-
-        text_embeddings = self.adapter.compute_text_embeddings(t, config, pooled_embeds)
-        image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
-            prompt_embeds,
-            config,
-            encoder_hidden_states_mask=encoder_mask,
-            cond_image_grid=cond_image_grid,
-            kontext_image_ids=kontext_image_ids,
-        )
-
         prev_patch_latents = [p for p in patch_latents]
+
         encoder_hidden_states: mx.array | None = None
 
         for patch_idx in range(len(patch_latents)):
@@ -809,31 +899,52 @@ class DiffusionRunner:
                 and not is_first_async_step
             ):
                 patch = mx.distributed.recv_like(
-                    patch, src=self.prev_rank, group=self.group
+                    patch, src=self.last_pipeline_rank, group=self.group
                 )
                 mx.eval(patch)
 
-            step_patch = mx.concatenate([patch, patch], axis=0) if needs_cfg else patch
+            results: list[tuple[bool, mx.array]] = []
 
-            noise, encoder_hidden_states = self._run_single_patch_pass(
-                patch=step_patch,
-                patch_idx=patch_idx,
-                token_indices=token_indices[patch_idx],
-                prompt_embeds=prompt_embeds,
-                text_embeddings=text_embeddings,
-                image_rotary_embeddings=image_rotary_embeddings,
-                encoder_hidden_states=encoder_hidden_states,
-            )
+            for branch in self._get_cfg_branches(prompt_data):
+                pooled_embeds = (
+                    branch.pooled if branch.pooled is not None else branch.embeds
+                )
+
+                text_seq_len = branch.embeds.shape[1]
+                self._ensure_wrappers(text_seq_len, branch.mask)
+                self._set_text_seq_len(text_seq_len)
+
+                if self.joint_block_wrappers:
+                    for wrapper in self.joint_block_wrappers:
+                        wrapper.set_encoder_mask(branch.mask)
+
+                text_embeddings = self.adapter.compute_text_embeddings(
+                    t, config, pooled_embeds
+                )
+                image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
+                    branch.embeds,
+                    config,
+                    encoder_hidden_states_mask=branch.mask,
+                    cond_image_grid=cond_image_grid,
+                    kontext_image_ids=kontext_image_ids,
+                )
+
+                noise, encoder_hidden_states = self._run_single_patch_pass(
+                    patch=patch,
+                    patch_idx=patch_idx,
+                    token_indices=token_indices[patch_idx],
+                    prompt_embeds=branch.embeds,
+                    text_embeddings=text_embeddings,
+                    image_rotary_embeddings=image_rotary_embeddings,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
+
+                if self.is_last_stage:
+                    assert noise is not None
+                    results.append((branch.positive, noise))
 
             if self.is_last_stage:
-                assert noise is not None
-                if needs_cfg:
-                    noise_pos, noise_neg = mx.split(noise, 2, axis=0)
-                    guidance_scale = self._get_effective_guidance_scale()
-                    assert guidance_scale is not None
-                    noise = self.adapter.apply_guidance(
-                        noise_pos, noise_neg, guidance_scale
-                    )
+                noise = self._combine_cfg_results(results)
 
                 patch_latents[patch_idx] = config.scheduler.step(  # pyright: ignore[reportAny]
                     noise=noise,
@@ -843,7 +954,9 @@ class DiffusionRunner:
 
                 if not self.is_first_stage and t != config.num_inference_steps - 1:
                     patch_latents[patch_idx] = mx.distributed.send(
-                        patch_latents[patch_idx], self.next_rank, group=self.group
+                        patch_latents[patch_idx],
+                        self.first_pipeline_rank,
+                        group=self.group,
                     )
                     mx.async_eval(patch_latents[patch_idx])
 
@@ -883,11 +996,12 @@ class DiffusionRunner:
 
         if self.has_joint_blocks:
             if not self.is_first_stage:
+                assert self.prev_pipeline_rank is not None
                 patch_len = patch.shape[1]
                 patch = mx.distributed.recv(
                     (batch_size, patch_len, hidden_dim),
                     patch.dtype,
-                    self.prev_rank,
+                    self.prev_pipeline_rank,
                     group=self.group,
                 )
                 mx.eval(patch)
@@ -896,7 +1010,7 @@ class DiffusionRunner:
                     encoder_hidden_states = mx.distributed.recv(
                         (batch_size, text_seq_len, hidden_dim),
                         patch.dtype,
-                        self.prev_rank,
+                        self.prev_pipeline_rank,
                         group=self.group,
                     )
                     mx.eval(encoder_hidden_states)
@@ -924,29 +1038,34 @@ class DiffusionRunner:
             if self.has_single_blocks or self.is_last_stage:
                 patch = patch_concat
             else:
+                assert self.next_pipeline_rank is not None
                 patch_concat = mx.distributed.send(
-                    patch_concat, self.next_rank, group=self.group
+                    patch_concat, self.next_pipeline_rank, group=self.group
                 )
                 mx.async_eval(patch_concat)
 
         elif self.has_joint_blocks and not self.is_last_stage:
-            patch = mx.distributed.send(patch, self.next_rank, group=self.group)
+            assert self.next_pipeline_rank is not None
+            patch = mx.distributed.send(
+                patch, self.next_pipeline_rank, group=self.group
+            )
             mx.async_eval(patch)
 
             if patch_idx == 0:
                 assert encoder_hidden_states is not None
                 encoder_hidden_states = mx.distributed.send(
-                    encoder_hidden_states, self.next_rank, group=self.group
+                    encoder_hidden_states, self.next_pipeline_rank, group=self.group
                 )
                 mx.async_eval(encoder_hidden_states)
 
         if self.has_single_blocks:
             if not self.owns_concat_stage and not self.is_first_stage:
+                assert self.prev_pipeline_rank is not None
                 patch_len = patch.shape[1]
                 patch = mx.distributed.recv(
                     (batch_size, text_seq_len + patch_len, hidden_dim),
                     patch.dtype,
-                    self.prev_rank,
+                    self.prev_pipeline_rank,
                     group=self.group,
                 )
                 mx.eval(patch)
@@ -961,7 +1080,10 @@ class DiffusionRunner:
                 )
 
             if not self.is_last_stage:
-                patch = mx.distributed.send(patch, self.next_rank, group=self.group)
+                assert self.next_pipeline_rank is not None
+                patch = mx.distributed.send(
+                    patch, self.next_pipeline_rank, group=self.group
+                )
                 mx.async_eval(patch)
 
         noise: mx.array | None = None

src/exo/worker/engines/mlx/auto_parallel.py

@@ -201,6 +201,9 @@ def pipeline_auto_parallel(
     device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size
 
     layers = layers[start_layer:end_layer]
+    for layer in layers:
+        mx.eval(layer)  # type: ignore
+
     layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
     layers[-1] = PipelineLastLayer(
         layers[-1],

src/exo/worker/runner/runner.py

@@ -61,7 +61,7 @@ from exo.shared.types.worker.runners import (
     RunnerStatus,
     RunnerWarmingUp,
 )
-from exo.shared.types.worker.shards import ShardMetadata
+from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 from exo.utils.channels import MpReceiver, MpSender
 from exo.worker.engines.image import (
     DistributedImageModel,
@@ -360,8 +360,9 @@ def main(
                         image_index = 0
                         for response in generate_image(model=model, task=task_params):
                             if (
-                                shard_metadata.device_rank
-                                == shard_metadata.world_size - 1
+                                isinstance(shard_metadata, PipelineShardMetadata)
+                                and shard_metadata.is_pipeline_last
+                                and shard_metadata.cfg_rank == 0
                             ):
                                 match response:
                                     case PartialImageResponse():
@@ -387,7 +388,11 @@ def main(
                                         image_index += 1
                     # can we make this more explicit?
                     except Exception as e:
-                        if shard_metadata.device_rank == shard_metadata.world_size - 1:
+                        if (
+                            isinstance(shard_metadata, PipelineShardMetadata)
+                            and shard_metadata.is_pipeline_last
+                            and shard_metadata.cfg_rank == 0
+                        ):
                             event_sender.send(
                                 ChunkGenerated(
                                     command_id=command_id,
@@ -419,8 +424,9 @@ def main(
                         image_index = 0
                         for response in generate_image(model=model, task=task_params):
                             if (
-                                shard_metadata.device_rank
-                                == shard_metadata.world_size - 1
+                                isinstance(shard_metadata, PipelineShardMetadata)
+                                and shard_metadata.is_pipeline_last
+                                and shard_metadata.cfg_rank == 0
                             ):
                                 match response:
                                     case PartialImageResponse():
@@ -445,7 +451,11 @@ def main(
                                         )
                                         image_index += 1
                     except Exception as e:
-                        if shard_metadata.device_rank == shard_metadata.world_size - 1:
+                        if (
+                            isinstance(shard_metadata, PipelineShardMetadata)
+                            and shard_metadata.is_pipeline_last
+                            and shard_metadata.cfg_rank == 0
+                        ):
                             event_sender.send(
                                 ChunkGenerated(
                                     command_id=command_id,