From 959a9de8d4904e6be2af65552ccf005b7086c066 Mon Sep 17 00:00:00 2001
From: justnodejs <justnodejs@gmail.com>
Date: Sat, 11 Apr 2026 17:02:18 +0800
Subject: [PATCH] fix(baoyu-imagine): tighten Replicate Seedream and Wan
 validation

---
 .../scripts/providers/replicate.test.ts       | 31 ++++++++++++++++++-
 .../scripts/providers/replicate.ts            | 30 ++++++------------
 2 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/skills/baoyu-imagine/scripts/providers/replicate.test.ts b/skills/baoyu-imagine/scripts/providers/replicate.test.ts
index c98486e..52cd825 100644
--- a/skills/baoyu-imagine/scripts/providers/replicate.test.ts
+++ b/skills/baoyu-imagine/scripts/providers/replicate.test.ts
@@ -182,6 +182,22 @@ test("Replicate input builder maps Wan models to their native schema", () => {
       thinking_mode: true,
     },
   );
+
+  assert.deepEqual(
+    buildInput(
+      "A robot painter",
+      "wan-video/wan-2.7-image",
+      makeArgs({
+        size: "1536x1024",
+      }),
+      [],
+    ),
+    {
+      prompt: "A robot painter",
+      size: "1536*1024",
+      thinking_mode: true,
+    },
+  );
 });
 
 test("Replicate input builder falls back to nano-banana schema for unknown models", () => {
@@ -245,9 +261,22 @@ test("Replicate validation catches unsupported Seedream and Wan argument combina
     /Wan image models on Replicate require --size/,
   );
 
+  assert.doesNotThrow(
+    () => validateArgs("wan-video/wan-2.7-image", makeArgs({ size: "1536x1024" })),
+  );
+
+  assert.doesNotThrow(
+    () => validateArgs("wan-video/wan-2.7-image-pro", makeArgs({ size: "1920x1080" })),
+  );
+
+  assert.throws(
+    () => validateArgs("wan-video/wan-2.7-image-pro", makeArgs({ referenceImages: Array.from({ length: 10 }, () => "ref.png") })),
+    /support at most 9 reference images/,
+  );
+
   assert.throws(
     () => validateArgs("wan-video/wan-2.7-image", makeArgs({ size: "4K" })),
-    /Wan image models on Replicate require --size to be one of/,
+    /Wan image models on Replicate require --size to be one of 1K, 2K or custom dimensions/,
   );
 
   assert.throws(
diff --git a/skills/baoyu-imagine/scripts/providers/replicate.ts b/skills/baoyu-imagine/scripts/providers/replicate.ts
index b07becc..2d6dd83 100644
--- a/skills/baoyu-imagine/scripts/providers/replicate.ts
+++ b/skills/baoyu-imagine/scripts/providers/replicate.ts
@@ -9,24 +9,8 @@ const MAX_POLL_MS = 300_000;
 const SIZE_PRESET_PATTERN = /^\d+K$/i;
 const SEEDREAM_45_SIZES = new Set(["2K", "4K"]);
 const SEEDREAM_5_LITE_SIZES = new Set(["2K", "3K"]);
-const WAN_PRO_SIZES = new Set([
-  "1K", "2K", "4K",
-  "1024*1024", "2048*2048", "4096*4096",
-  "1280*720", "720*1280",
-  "2048*1152", "1152*2048",
-  "4096*2304", "2304*4096",
-  "1024*768", "768*1024",
-  "2048*1536", "1536*2048",
-  "4096*3072", "3072*4096",
-]);
-const WAN_SIZES = new Set([
-  "1K", "2K",
-  "1024*1024", "2048*2048",
-  "1280*720", "720*1280",
-  "2048*1152", "1152*2048",
-  "1024*768", "768*1024",
-  "2048*1536", "1536*2048",
-]);
+const WAN_PRO_PRESET_SIZES = new Set(["1K", "2K", "4K"]);
+const WAN_PRESET_SIZES = new Set(["1K", "2K"]);
 
 export function getDefaultModel(): string {
   return process.env.REPLICATE_IMAGE_MODEL || DEFAULT_MODEL;
@@ -116,7 +100,7 @@ function getAllowedSeedreamSizes(model: string): Set<string> {
 }
 
 function getAllowedWanSizes(model: string): Set<string> {
-  return isWanProModel(model) ? WAN_PRO_SIZES : WAN_SIZES;
+  return isWanProModel(model) ? WAN_PRO_PRESET_SIZES : WAN_PRESET_SIZES;
 }
 
 function normalizePresetSize(size: string): string {
@@ -265,11 +249,15 @@ export function validateArgs(model: string, args: CliArgs): void {
       throw new Error("Wan image models on Replicate require --size when using --ar. This provider does not infer size from aspect ratio.");
     }
 
+    if (args.referenceImages.length > 9) {
+      throw new Error("Wan image models on Replicate support at most 9 reference images per request.");
+    }
+
     if (args.size) {
       const normalizedSize = parsePixelSize(args.size) ? normalizePixelSize(args.size) : normalizePresetSize(args.size);
-      if (!getAllowedWanSizes(model).has(normalizedSize)) {
+      if (!parsePixelSize(args.size) && !getAllowedWanSizes(model).has(normalizedSize)) {
         throw new Error(
-          `Wan image models on Replicate require --size to be one of ${Array.from(getAllowedWanSizes(model)).join(", ")}. Received: ${args.size}`
+          `Wan image models on Replicate require --size to be one of ${Array.from(getAllowedWanSizes(model)).join(", ")} or custom dimensions like 1920x1080. Received: ${args.size}`
         );
       }
     }