asset_train_brand_lora

Instructions

Implementation Reference

• packages/mcp-server/src/tools/train-brand-lora.ts:37-121 (handler)

  The trainBrandLora function is the core handler for the asset_train_brand_lora tool. It reads training images from local filesystem paths, validates them (15-100 count warnings), converts to base64, POSTs to a user-owned LoRA training endpoint (PROMPT_TO_BUNDLE_MODAL_LORA_TRAIN_URL), and returns the lora_id, status, and lora_url on success.

  export async function trainBrandLora(input: TrainBrandLoraInputT): Promise<{
    ok: boolean;
    lora_id?: string;
    status?: "ready" | "training";
    lora_url?: string;
    error?: string;
    warnings: string[];
  }> {
    const url = process.env["PROMPT_TO_BUNDLE_MODAL_LORA_TRAIN_URL"];
    if (!url) {
      return {
        ok: false,
        error:
          "PROMPT_TO_BUNDLE_MODAL_LORA_TRAIN_URL not set. Point it at a LoRA-training endpoint. Reference implementations: Modal + ai-toolkit (https://modal.com/docs/examples/flux_lora), Replicate replicate/flux-dev-lora-trainer. See docs/research/06-stable-diffusion-flux/6d-lora-training-for-brand-style.md.",
        warnings: []
      };
    }
    const token = process.env["PROMPT_TO_BUNDLE_MODAL_LORA_TRAIN_TOKEN"];
  
    const warnings: string[] = [];
    if (input.training_images.length < 15) {
      warnings.push(
        `only ${input.training_images.length} training images provided — LoRA quality degrades below ~20. Research suggests 20-50 curated shots.`
      );
    }
    if (input.training_images.length > 100) {
      warnings.push(
        `${input.training_images.length} training images — training cost scales linearly; consider curating to <=50 representative shots.`
      );
    }
  
    // Each path goes through safeReadPath so a crafted MCP input can't read
    // outside the project's allow-list. Base64 payload stays local until POST.
    const imagesBase64: string[] = [];
    for (const p of input.training_images) {
      const buf = readFileSync(safeReadPath(p));
      imagesBase64.push(buf.toString("base64"));
    }
  
    const body: Record<string, unknown> = {
      name: input.name,
      base_model: input.base_model,
      training_images: imagesBase64,
      rank: input.rank ?? 16,
      steps: input.steps ?? 1200
    };
    if (input.captions && input.captions.length > 0) body["captions"] = input.captions;
  
    let resp: Response;
    try {
      resp = await fetch(url, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          ...(token && { Authorization: `Bearer ${token}` })
        },
        body: JSON.stringify(body),
        signal: AbortSignal.timeout(600_000)
      });
    } catch (err) {
      return { ok: false, error: `fetch failed: ${(err as Error).message}`, warnings };
    }
  
    if (!resp.ok) {
      const errText = await resp.text();
      return { ok: false, error: `HTTP ${resp.status}: ${errText.slice(0, 400)}`, warnings };
    }
  
    const json = (await resp.json()) as {
      lora_id?: string;
      status?: "ready" | "training";
      lora_url?: string;
    };
    if (!json.lora_id) {
      return { ok: false, error: "training endpoint did not return lora_id", warnings };
    }
  
    return {
      ok: true,
      lora_id: json.lora_id,
      status: json.status ?? "training",
      ...(json.lora_url && { lora_url: json.lora_url }),
      warnings
    };
  }

• packages/mcp-server/src/schemas.ts:207-228 (schema)

  Zod schema (TrainBrandLoraInput) defining the input validation: name (string, min 1), base_model (default sdxl-1.0), training_images (array of strings, 5-200 items), captions (optional array of strings), rank (4-128, default 16), steps (100-8000, default 1200).

  export const TrainBrandLoraInput = z.object({
    name: z.string().min(1).describe("Brand slug. Becomes the LoRA id trigger token."),
    base_model: z
      .string()
      .default("sdxl-1.0")
      .describe(
        "Base model to fine-tune. Supported by typical trainers: sdxl-1.0, flux-1-dev, sd-1.5."
      ),
    training_images: z
      .array(z.string())
      .min(5)
      .max(200)
      .describe(
        "Local filesystem paths to 5-200 brand-consistent images. 20-50 is the sweet spot. Paths go through the same safeReadPath allow-list as the other tools."
      ),
    captions: z
      .array(z.string())
      .optional()
      .describe("Per-image caption overrides. If omitted, the trainer auto-captions."),
    rank: z.number().int().min(4).max(128).default(16).describe("LoRA rank."),
    steps: z.number().int().min(100).max(8000).default(1200).describe("Training steps.")
  });

• packages/mcp-server/src/server.ts:491-521 (registration)

  Tool registration in the TOOLS array: defines name 'asset_train_brand_lora', description, inputSchema with properties (name, base_model, training_images, captions, rank, steps), required fields, and annotations.

  {
    name: "asset_train_brand_lora",
    description:
      "Train a brand-consistent LoRA from 20-50 sample images, returning a `lora_id` the `comfyui-*` and SDXL-family providers can reference. Requires a user-owned training endpoint (Modal / Runpod / self-host) at PROMPT_TO_BUNDLE_MODAL_LORA_TRAIN_URL. Phase-4 scaffold: the MCP tool does the packaging, validation, and HTTP; the user owns the deployment and pricing. See docs/research/06-stable-diffusion-flux/6d-lora-training-for-brand-style.md.",
    inputSchema: {
      type: "object",
      properties: {
        name: { type: "string", description: "Brand slug. Becomes the LoRA trigger token." },
        base_model: {
          type: "string",
          description: "Base model to fine-tune (sdxl-1.0 / flux-1-dev / sd-1.5).",
          default: "sdxl-1.0"
        },
        training_images: {
          type: "array",
          items: { type: "string" },
          description:
            "Local filesystem paths (5-200). 20-50 is the sweet spot. Paths go through the safeReadPath allow-list."
        },
        captions: {
          type: "array",
          items: { type: "string" },
          description: "Per-image caption overrides. Auto-captioned if omitted."
        },
        rank: { type: "number", default: 16 },
        steps: { type: "number", default: 1200 }
      },
      required: ["name", "training_images"]
    },
    annotations: { openWorldHint: true }
  },

• packages/mcp-server/src/server.ts:831-833 (registration)

  Handler dispatch in the CallToolRequestSchema switch statement: case 'asset_train_brand_lora' calls trainBrandLora(TrainBrandLoraInput.parse(args ?? {})).

  case "asset_train_brand_lora":
    result = await trainBrandLora(TrainBrandLoraInput.parse(args ?? {}));
    break;

• packages/mcp-server/src/server.ts:54-54 (helper)

  Import binding of the trainBrandLora function from './tools/train-brand-lora.js' into the server module.

  import { trainBrandLora } from "./tools/train-brand-lora.js";