Move the ui model cards to their own plugin folders so plugins can add custom models in the trainer/generator/ui

extensions_built_in/audio_models/ui.tsx

@@ -0,0 +1,268 @@
+// UI entries (training form + Generate page) for the models this package
+// registers in AI_TOOLKIT_MODELS. Loaded at runtime by the UI, not bundled:
+// see ui/src/extensions/README.md for the convention and the allowed imports.
+import type { ModelArch } from "@/app/jobs/new/options";
+import {
+  defaultSampleConfig,
+  defaultAudioSampleConfig,
+  defaultYue2SampleConfig,
+} from "@/helpers/defaultSamples";
+
+const defaultNameOrPath = "";
+
+export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
+  {
+    name: "ace_step_15_xl",
+    label: "ACE-Step 1.5 XL",
+    group: "audio",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "config.process[0].model.name_or_path": [
+        "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_xl_base_aio.safetensors",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.low_vram": [true, false],
+      "config.process[0].train.unload_text_encoder": [false, false],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.timestep_type": ["linear", "sigmoid"],
+      "config.process[0].model.qtype": ["qfloat8", "qfloat8"],
+      "config.process[0].sample": [
+        defaultAudioSampleConfig,
+        defaultSampleConfig,
+      ],
+    },
+    sampleTags: {
+      CAPTION: {
+        title: "Audio Prompt",
+        type: "text",
+        full: true,
+      },
+      LYRICS: {
+        title: "Lyrics",
+        type: "multiline",
+        full: true,
+      },
+      BPM: {
+        title: "BPM",
+        type: "number",
+      },
+      KEYSCALE: {
+        title: "Key Scale",
+        type: "text",
+      },
+      TIMESIGNATURE: {
+        title: "Time Signature",
+        type: "text",
+      },
+      DURATION: {
+        title: "Duration (sec)",
+        type: "number",
+      },
+      LANGUAGE: {
+        title: "Language",
+        type: "text",
+      },
+    },
+    disableSections: ["network.conv"],
+    additionalSections: [
+      "sample.multi_ctrl_imgs",
+      "model.low_vram",
+      "model.layer_offloading",
+    ],
+  },
+  {
+    name: "yue2",
+    label: "YuE2",
+    group: "audio",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "config.process[0].model.name_or_path": [
+        "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.quantize_te": [false, false],
+      "config.process[0].model.low_vram": [false, false],
+      "config.process[0].train.unload_text_encoder": [false, false],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.timestep_type": ["sigmoid", "sigmoid"],
+      // the int8 repack ships convrot8 layers; requesting convrot8 keeps them as-is (no requantization)
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      "config.process[0].sample": [
+        defaultYue2SampleConfig,
+        defaultSampleConfig,
+      ],
+      "config.process[0].datasets[x].cache_latents_to_disk": [true, true],
+      // audio has no resolution; every bucket would duplicate the whole dataset
+      "config.process[0].datasets[x].resolution": [[512], [512, 768, 1024]],
+      // blank captions break lyric following; the AR must always see the prefix
+      "config.process[0].datasets[x].caption_dropout_rate": [0, 0.05],
+      "config.process[0].model.model_kwargs": [
+        {
+          cot: "full",
+          abc_dropout: 0.5,
+          sample_ar_repetition_penalty: 1.2,
+          ar_kl_weight: 0.2,
+        },
+        {},
+      ],
+    },
+    // native YuE2 prompt: style text, a [Lyrics] line, the lyrics
+    hasMultiLinePrompts: true,
+    modelNotes: (
+      <div className="space-y-2">
+        <p className="font-semibold text-amber-400">
+          Experimental. The AR (composition) model memorizes quickly and does
+          not work well on small datasets. A handful of songs is enough for it
+          to learn the exact token sequence of each song; after that it stops
+          generalizing and free-running samples drift away from the training
+          material. Expect to need a large, varied dataset for the AR side to
+          learn a style rather than the songs themselves.
+        </p>
+        <p>
+          YuE2 is two experts on one backbone. The <b>AR expert</b> reads the
+          style line and lyrics and writes the song as a sequence of semantic
+          codec tokens (25 per second). The <b>NAR expert</b> then renders those
+          tokens into audio latents with flow matching. Training a LoRA here
+          trains both: next-token loss on the AR over the whole song from its
+          start, flow loss on the NAR over a random window.
+        </p>
+        <p>
+          The AR loss is the one to watch (<code>loss/ar_ce</code>). It starts
+          near 5 and, on a small dataset, falls toward 0 within a few hundred
+          steps, which is memorization. <code>ar_kl_weight</code> in model
+          kwargs anchors the AR to the base model so it cannot collapse onto the
+          training songs; <code>ar_lr_multiplier</code> and a smaller AR rank
+          slow it further. Style comes mostly from the NAR, lyric following from
+          the AR. Do not use caption dropout: a blank prompt breaks lyric
+          following.
+        </p>
+        <p>
+          The official audio-to-token encoder is unreleased. Training uses the
+          community tokenizer by Kytra (
+          <a
+            href="https://x.com/sin_ceriously"
+            target="_blank"
+            rel="noreferrer"
+            className="text-blue-400 hover:underline"
+          >
+            @sin_ceriously
+          </a>
+          ), a MERT-v2-FullSong head that maps real audio to YuE2 codec tokens:{" "}
+          <a
+            href="https://huggingface.co/Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4"
+            target="_blank"
+            rel="noreferrer"
+            className="text-blue-400 hover:underline"
+          >
+            Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4
+          </a>
+          . It is downloaded on first use. Its tokens are an approximation of
+          the model's own, so rendered samples will not be bit-faithful to the
+          training audio even when the AR replays a song exactly. Though it is
+          extremely close.
+        </p>
+        <p>
+          <b>ABC generation.</b> Generation is two stages: the AR first writes a
+          lead sheet of the whole song in ABC notation (sections, chords, vocal
+          and instrumental melody), then writes the codec tokens conditioned on
+          that sheet. It can also run without a sheet ("off" mode), which uses a
+          different instruction line. Samples here do the two stages with{" "}
+          <code>cot: full</code> (chords) or <code>cot: melody</code> (melody
+          only), or the single stage with <code>cot: off</code>.
+        </p>
+        <p>
+          <b>Training for both.</b> Every song is transcribed to an ABC sheet
+          with SheetSage2 and the AR is trained on lyrics to sheet and sheet to
+          tokens. <code>abc_dropout</code> (default 0.5) is the fraction of
+          training items fed without the sheet instead, as an off-mode prompt,
+          so one LoRA works in both modes. Set it to 0 to train the sheet path
+          only, or 1 to train off mode only.
+        </p>
+        <p>
+          <b>Caching is required.</b> The sheet is produced at latent-cache time
+          and stored with the latents and codec tokens, so Cache Latents to Disk
+          must stay on: without it, SheetSage2 (about 12 s per song), the MERT
+          tokenizer and the VAE would run again on every training step. The
+          cache records which mode built it; changing <code>cot</code> means
+          deleting the dataset's <code>_latent_cache</code> folder so the sheets
+          are rebuilt. Audio has no resolution, so keep a single resolution
+          bucket per dataset or every bucket duplicates the songs.
+        </p>
+        <p>
+          Prompt format: a style line, then <code>[Lyrics]</code>, then the
+          lyrics with bracketed section headers such as <code>[Verse 1]</code>{" "}
+          and <code>[Chorus]</code>. The Qwen3-Omni captioner has a YuE2 preset
+          that writes captions in this layout. Sample length is set by the
+          Duration field in the sample section.
+        </p>
+      </div>
+    ),
+    // no separate text encoder: the prompt side is the AR expert, covered by the transformer quantization
+    disableSections: ["network.conv", "model.quantize_te"],
+    additionalSections: ["model.low_vram", "sample.duration"],
+  },
+  {
+    name: "ace_step_15",
+    label: "ACE-Step 1.5",
+    group: "audio",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "config.process[0].model.name_or_path": [
+        "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.low_vram": [true, false],
+      "config.process[0].train.unload_text_encoder": [false, false],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.timestep_type": ["linear", "sigmoid"],
+      "config.process[0].model.qtype": ["qfloat8", "qfloat8"],
+      "config.process[0].sample": [
+        defaultAudioSampleConfig,
+        defaultSampleConfig,
+      ],
+    },
+    sampleTags: {
+      CAPTION: {
+        title: "Audio Prompt",
+        type: "text",
+        full: true,
+      },
+      LYRICS: {
+        title: "Lyrics",
+        type: "multiline",
+        full: true,
+      },
+      BPM: {
+        title: "BPM",
+        type: "number",
+      },
+      KEYSCALE: {
+        title: "Key Scale",
+        type: "text",
+      },
+      TIMESIGNATURE: {
+        title: "Time Signature",
+        type: "text",
+      },
+      DURATION: {
+        title: "Duration (sec)",
+        type: "number",
+      },
+      LANGUAGE: {
+        title: "Language",
+        type: "text",
+      },
+    },
+    disableSections: ["network.conv"],
+    additionalSections: [
+      "sample.multi_ctrl_imgs",
+      "model.low_vram",
+      "model.layer_offloading",
+    ],
+  },
+];

extensions_built_in/diffusion_models/example_model/README.md

@@ -34,9 +34,10 @@ example/
    own `AI_TOOLKIT_MODELS` list — see `extensions/z_image_pixel/`.)
 3. The class attribute `arch` (e.g. `"example"`) is matched against
    `model.arch` in the training config YAML to pick your class.
-4. To expose it in the web UI, add an entry to
+4. To expose it in the web UI, add an entry to your package's `ui.tsx`
-   `ui/src/app/jobs/new/options.ts` (search for an existing arch like
+   (`AI_TOOLKIT_UI_MODELS`, see `ui/src/extensions/README.md`; built-in
-   `ideogram4` to copy the shape).
+   models live in `extensions_built_in/diffusion_models/ui.tsx`). It is
+   loaded at runtime, so no UI rebuild is needed.
 
 Minimal config YAML to train it:
 

extensions_built_in/flex2/ui.tsx

@@ -0,0 +1,40 @@
+// UI entries (training form + Generate page) for the models this package
+// registers in AI_TOOLKIT_MODELS. Loaded at runtime by the UI, not bundled:
+// see ui/src/extensions/README.md for the convention and the allowed imports.
+import type { ModelArch } from "@/app/jobs/new/options";
+
+const defaultNameOrPath = "";
+
+export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
+  {
+    name: "flex2",
+    label: "Flex.2",
+    group: "image",
+    controls: ["depth", "line", "pose", "inpaint"],
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "config.process[0].model.name_or_path": [
+        "ostris/Flex.2-preview",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.model_kwargs": [
+        {
+          invert_inpaint_mask_chance: 0.2,
+          inpaint_dropout: 0.5,
+          control_dropout: 0.5,
+          inpaint_random_chance: 0.2,
+          do_random_inpainting: true,
+          random_blur_mask: true,
+          random_dialate_mask: true,
+        },
+        {},
+      ],
+      "config.process[0].train.bypass_guidance_embedding": [true, false],
+      "config.process[0].sample.sampler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+    },
+    disableSections: ["network.conv"],
+  },
+];

extensions_built_in/llm_models/ui.tsx

@@ -0,0 +1,84 @@
+// UI entries (training form + Generate page) for the models this package
+// registers in AI_TOOLKIT_MODELS. Loaded at runtime by the UI, not bundled:
+// see ui/src/extensions/README.md for the convention and the allowed imports.
+import type { ModelArch } from "@/app/jobs/new/options";
+import {
+  defaultSampleConfig,
+  defaultQwen25OmniSampleConfig,
+} from "@/helpers/defaultSamples";
+
+const defaultNameOrPath = "";
+
+export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
+  {
+    name: "qwen25_omni",
+    label: "Qwen2.5-Omni",
+    group: "llm",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "config.process[0].model.name_or_path": [
+        "ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.quantize_te": [false, false],
+      "config.process[0].model.low_vram": [false, false],
+      // the single-file thinker ships convrot8 layers; requesting convrot8 keeps them as-is
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      "config.process[0].train.unload_text_encoder": [false, false],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.batch_size": [1, 1],
+      "config.process[0].sample": [
+        defaultQwen25OmniSampleConfig,
+        defaultSampleConfig,
+      ],
+      // media is encoded on the GPU per step; the cache is optional and large (~54 MB per 300 s of audio)
+      "config.process[0].datasets[x].cache_latents_to_disk": [false, true],
+      "config.process[0].datasets[x].resolution": [[512], [512, 768, 1024]],
+      // the caption is the training target; a blank one trains nothing
+      "config.process[0].datasets[x].caption_dropout_rate": [0, 0.05],
+      "config.process[0].model.model_kwargs": [
+        { instruction: "Describe this in detail." },
+        {},
+      ],
+    },
+    disableSections: [
+      "network.conv",
+      "trigger_word",
+      "train.diff_output_preservation",
+      "train.blank_prompt_preservation",
+      "train.unload_text_encoder",
+      "slider",
+    ],
+    additionalSections: [
+      "model.model_kwargs.instruction",
+      "sample.ctrl_img",
+      "datasets.num_frames",
+    ],
+    modelNotes: (
+      <div className="space-y-2">
+        <p>
+          Qwen2.5-Omni 7B thinker as a text-generating model: audio, image or
+          video in, text out. Each dataset item is a media file (mp3, wav, flac,
+          ogg, jpg, png, webp, mp4, ...) with a caption file next to it; the
+          caption is the text the model learns to produce for that media. One
+          dataset can mix all three kinds. Video files are used when Num Frames
+          is above 1 and are seen as frames only (no audio track).
+        </p>
+        <p>
+          The instruction in <code>model_kwargs.instruction</code> is the user
+          turn for every training item; use the same wording when captioning
+          with the trained LoRA. Samples take a media file per prompt and write
+          the generated text as a .txt file.
+        </p>
+        <p>
+          Media is encoded by the frozen audio and vision towers on the GPU each
+          step, so Cache Latents to Disk can stay off (the cache is about 54 MB
+          per 300 seconds of audio). Batch size is 1. The LoRA trains the text
+          stack only. Watch <code>loss/ce</code>: next-token cross-entropy on
+          the caption.
+        </p>
+      </div>
+    ),
+  },
+];

tsconfig.json

@@ -0,0 +1,28 @@
+{
+  // type checking + editor support for extension UI modules only (the web
+  // app has its own config in ui/); they are loaded at runtime, not bundled
+  "extends": "./ui/tsconfig.json",
+  "compilerOptions": {
+    "noEmit": true,
+    "incremental": false,
+    "plugins": [],
+    "typeRoots": [
+      "./ui/node_modules/@types"
+    ],
+    "paths": {
+      "@/*": [
+        "./ui/src/*"
+      ],
+      "*": [
+        "./ui/node_modules/@types/*",
+        "./ui/node_modules/*"
+      ]
+    }
+  },
+  "include": [
+    "extensions/*/ui.tsx",
+    "extensions/*/ui.ts",
+    "extensions_built_in/*/ui.tsx",
+    "extensions_built_in/*/ui.ts"
+  ]
+}

ui/package.json

@@ -11,7 +11,8 @@
     "build_and_start": "npm run install_deps && npm run db_build_start",
     "lint": "next lint",
     "update_db": "npx prisma generate && npx prisma db push",
-    "format": "prettier --write \"**/*.{js,jsx,ts,tsx,css,scss}\""
+    "format": "prettier --write \"**/*.{js,jsx,ts,tsx,css,scss}\"",
+    "check_extensions": "tsc -p ../tsconfig.json"
   },
   "dependencies": {
     "@headlessui/react": "^2.2.0",

ui/src/app/api/model_archs/route.ts

@@ -0,0 +1,17 @@
+import { NextResponse } from 'next/server';
+import { listExtensionUiModules } from '@/server/extensionUi';
+
+// Extension UI modules (AI_TOOLKIT_UI_MODELS) as transpiled CommonJS source,
+// evaluated client-side by useModelArchs. Read from disk per request so new
+// models appear without a rebuild.
+export const runtime = 'nodejs';
+export const dynamic = 'force-dynamic';
+
+export async function GET() {
+  try {
+    const result = await listExtensionUiModules();
+    return NextResponse.json(result, { headers: { 'Cache-Control': 'no-store' } });
+  } catch (e: any) {
+    return NextResponse.json({ error: e?.message || String(e) }, { status: 500 });
+  }
+}

ui/src/app/generate/page.tsx

@@ -36,7 +36,8 @@ import { defaultInferenceJobConfig } from '@/helpers/inferenceJobConfig';
 import { encodeFilePathForUrl, objToTags, tagsToObj } from '@/utils/basic';
 import { latentToImage, payloadToFloat32, readEngineFrames, PreviewInfo } from '@/utils/engineStream';
 import { isMac } from '@/helpers/basic';
-import { modelArchs, getGenerateDefaults, GenerateDefaults } from '@/app/jobs/new/options';
+import { getGenerateDefaults, GenerateDefaults } from '@/app/jobs/new/options';
+import { useModelArchs } from '@/extensions/modelArchs';
 import GenerateFooter from '@/components/generate/GenerateFooter';
 import LoraBrowserModal, { LoraPick } from '@/components/generate/LoraBrowserModal';
 
@@ -296,9 +297,10 @@ function GeneratePageInner() {
 
   // ---- models: the same arch list the training UI uses (jobs/new/options.tsx) ----
   // text-generating archs have no Generate page path yet
+  const { archs: modelArchs } = useModelArchs();
   const archs: GenerateDefaults[] = useMemo(
     () => modelArchs.filter(a => a.group !== 'llm').map(getGenerateDefaults),
-    [],
+    [modelArchs],
   );
   const ready = !!engineStatus?.running;
 

ui/src/app/jobs/new/SimpleJob.tsx

@@ -1,14 +1,13 @@
 'use client';
 import { useMemo } from 'react';
 import {
-  modelArchs,
   ModelArch,
-  groupedModelOptions,
   quantizationOptions,
   defaultQtype,
   jobTypeOptions,
   SampleTags,
 } from './options';
+import { useModelArchs } from '@/extensions/modelArchs';
 import { defaultCompileOptions, defaultDatasetConfig } from './jobConfig';
 import { GroupedSelectOption, JobConfig, SelectOption } from '@/types';
 import { objectCopy, tagsToObj, objToTags } from '@/utils/basic';
@@ -61,9 +60,10 @@ export default function SimpleJob({
   datasetOptions,
   isLoading,
 }: Props) {
+  const { archs: modelArchs, groupedModelOptions } = useModelArchs();
   const modelArch = useMemo(() => {
     return modelArchs.find(a => a.name === jobConfig.config.process[0].model.arch) as ModelArch;
-  }, [jobConfig.config.process[0].model.arch]);
+  }, [modelArchs, jobConfig.config.process[0].model.arch]);
 
   const jobType = useMemo(() => {
     return jobTypeOptions.find(j => j.value === jobConfig.config.process[0].type);
@@ -284,7 +284,7 @@ export default function SimpleJob({
               label="Model Architecture"
               value={jobConfig.config.process[0].model.arch}
               onChange={value => {
-                handleModelArchChange(jobConfig.config.process[0].model.arch, value, jobConfig, setJobConfig);
+                handleModelArchChange(modelArchs, jobConfig.config.process[0].model.arch, value, jobConfig, setJobConfig);
               }}
               options={groupedModelOptions}
             />

ui/src/app/jobs/new/utils.ts

@@ -1,5 +1,5 @@
 import { GroupedSelectOption, JobConfig, SelectOption } from '@/types';
-import { modelArchs, ModelArch } from './options';
+import { ModelArch } from './options';
 import { objectCopy } from '@/utils/basic';
 
 const expandDatasetDefaults = (
@@ -22,6 +22,7 @@ const expandDatasetDefaults = (
 };
 
 export const handleModelArchChange = (
+  modelArchs: ModelArch[],
   currentArchName: string,
   newArchName: string,
   jobConfig: JobConfig,

ui/src/extensions/README.md

@@ -0,0 +1,37 @@
+# Extension UI modules
+
+A model package registers its python classes in `AI_TOOLKIT_MODELS`
+(`extensions/<pkg>/__init__.py`). To show those archs in the web UI (training
+form and Generate page) add a `ui.tsx` next to it exporting the UI entries:
+
+```tsx
+// extensions/<pkg>/ui.tsx
+import type { ModelArch } from '@/app/jobs/new/options';
+
+export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
+  {
+    name: 'my_arch', // must equal the python class's `arch`
+    label: 'My Model',
+    group: 'image',
+    defaults: {
+      'config.process[0].model.name_or_path': ['org/my-model', ''],
+      'config.process[0].model.quantize': [true, false],
+    },
+  },
+];
+```
+
+- The file is read and transpiled by the UI server on every page load
+  (`/api/model_archs`), so adding or editing one needs no rebuild. `.ts`,
+  `.jsx` and `.js` also work.
+- The `ModelArch` shape is in `ui/src/app/jobs/new/options.tsx`; the built-in
+  tables in `extensions_built_in/*/ui.tsx` are the reference examples.
+- Code runs in the browser. Only these imports are available: `react`,
+  `next/link`, `@/helpers/defaultSamples`, `@/components/formInputs`, plus
+  `@/app/jobs/new/options`, and type-only imports from `@/types`.
+- Entries are merged by `name`; `extensions/` loads after
+  `extensions_built_in/`, so an extension can override a built-in entry.
+- A module that fails to compile or evaluate is skipped and logged in the
+  browser console; the others still load.
+- Type check: `cd ui && npm run check_extensions` (uses the repo root
+  `tsconfig.json`, which editors also pick up for these files).

ui/src/extensions/modelArchs.ts

@@ -0,0 +1,108 @@
+'use client';
+// Runtime-loaded model arch list. Extension packages ship a ui.tsx exporting
+// AI_TOOLKIT_UI_MODELS (see README.md here); /api/model_archs transpiles them
+// and this hook evaluates them in the browser with a fixed set of importable
+// modules. Later modules override earlier ones by arch name, so a package
+// under extensions/ can replace a built-in entry.
+import React, { useEffect, useMemo } from 'react';
+import * as ReactJsxRuntime from 'react/jsx-runtime';
+import * as ReactJsxDevRuntime from 'react/jsx-dev-runtime';
+import Link from 'next/link';
+import { createGlobalState } from 'react-global-hooks';
+import * as defaultSamples from '@/helpers/defaultSamples';
+import * as formInputs from '@/components/formInputs';
+import { apiClient } from '@/utils/api';
+import * as options from '@/app/jobs/new/options';
+import { ModelArch } from '@/app/jobs/new/options';
+import { GroupedSelectOption } from '@/types';
+
+// what `import ... from '<name>'` resolves to inside an extension ui module
+const SHIM: { [name: string]: any } = {
+  react: React,
+  'react/jsx-runtime': ReactJsxRuntime,
+  'react/jsx-dev-runtime': ReactJsxDevRuntime,
+  'next/link': { __esModule: true, default: Link },
+  '@/helpers/defaultSamples': defaultSamples,
+  '@/components/formInputs': formInputs,
+  '@/types': {}, // types only
+  '@/app/jobs/new/options': options,
+};
+
+interface UiModule {
+  id: string;
+  file: string;
+  mtimeMs: number;
+  code: string;
+}
+
+interface ModelArchsState {
+  archs: ModelArch[];
+  errors: string[];
+  isLoaded: boolean;
+}
+
+const state = createGlobalState<ModelArchsState>({ archs: [], errors: [], isLoaded: false });
+let inflight: Promise<void> | null = null;
+
+const evaluateModule = (mod: UiModule): ModelArch[] => {
+  const module = { exports: {} as any };
+  const require = (name: string) => {
+    if (name in SHIM) return SHIM[name];
+    throw new Error(`cannot import "${name}" (available: ${Object.keys(SHIM).join(', ')})`);
+  };
+  new Function('require', 'module', 'exports', mod.code)(require, module, module.exports);
+  const list = module.exports.AI_TOOLKIT_UI_MODELS;
+  if (!Array.isArray(list)) throw new Error('must export an AI_TOOLKIT_UI_MODELS array');
+  list.forEach((a, i) => {
+    if (!a || typeof a.name !== 'string' || typeof a.label !== 'string' || typeof a.group !== 'string') {
+      throw new Error(`AI_TOOLKIT_UI_MODELS[${i}] needs string name, label and group`);
+    }
+  });
+  return list as ModelArch[];
+};
+
+export const loadModelArchs = async (): Promise<void> => {
+  if (inflight) return inflight;
+  inflight = (async () => {
+    try {
+      const res = await apiClient.get('/api/model_archs');
+      const modules: UiModule[] = res.data?.modules || [];
+      const errors: string[] = (res.data?.errors || []).map((e: any) => `${e.id}: ${e.error}`);
+      const byName = new Map<string, ModelArch>();
+      for (const mod of modules) {
+        try {
+          for (const arch of evaluateModule(mod)) byName.set(arch.name, arch);
+        } catch (e: any) {
+          errors.push(`${mod.id}: ${e?.message || e}`);
+        }
+      }
+      for (const err of errors) console.error('[model_archs]', err);
+      const archs = [...byName.values()].sort((a, b) =>
+        a.label.localeCompare(b.label, undefined, { sensitivity: 'base' }),
+      );
+      state.set({ archs, errors, isLoaded: true });
+    } catch (e: any) {
+      const msg = e?.response?.data?.error || e?.message || String(e);
+      console.error('[model_archs]', msg);
+      state.set({ ...state.get(), errors: [msg], isLoaded: true });
+    } finally {
+      inflight = null;
+    }
+  })();
+  return inflight;
+};
+
+/** Model archs for the training form and Generate page. Refetches on every
+ * mount (server side is mtime-cached), so navigating to a page picks up new
+ * extension files; previous data stays until the reload lands. */
+export const useModelArchs = (): ModelArchsState & {
+  groupedModelOptions: GroupedSelectOption[];
+  refresh: () => Promise<void>;
+} => {
+  const [current] = state.use();
+  useEffect(() => {
+    loadModelArchs();
+  }, []);
+  const groupedModelOptions = useMemo(() => options.groupModelOptions(current.archs), [current.archs]);
+  return { ...current, groupedModelOptions, refresh: loadModelArchs };
+};

ui/src/server/extensionUi.ts

@@ -0,0 +1,122 @@
+// Discovers and transpiles extension UI modules at request time.
+//
+// Every package under extensions/ and extensions_built_in/ may ship a
+// `ui.tsx` (or .ts/.jsx/.js) exporting AI_TOOLKIT_UI_MODELS, the UI
+// counterpart of the python AI_TOOLKIT_MODELS list. The files are read from
+// disk on each call (mtime-cached), so a new or edited model shows up on the
+// next page load without a rebuild. Transpiled with the SWC that ships inside
+// Next (native binding, all platforms); evaluation happens in the browser,
+// see src/extensions/modelArchs.ts.
+import fs from 'fs';
+import path from 'path';
+import { TOOLKIT_ROOT } from '@/paths';
+
+const EXTENSION_DIRS = ['extensions_built_in', 'extensions'];
+const UI_FILE_NAMES = ['ui.tsx', 'ui.ts', 'ui.jsx', 'ui.js'];
+
+export interface ExtensionUiModule {
+  id: string; // "<dir>/<package>"
+  file: string; // absolute path
+  mtimeMs: number;
+  code: string; // CommonJS source
+}
+
+export interface ExtensionUiError {
+  id: string;
+  file: string;
+  error: string;
+}
+
+const cache = new Map<string, ExtensionUiModule>();
+
+// the swc binding is a native module resolved at runtime, outside the bundler
+// graph. webpack rewrites a statically imported createRequire whose base path
+// it cannot evaluate into `(void 0)(...)`, so the builtin is loaded through an
+// import both bundlers are told to leave alone.
+let swcPromise: Promise<any> | null = null;
+const getSwc = () => {
+  if (!swcPromise) {
+    swcPromise = (async () => {
+      const nodeModule = await import(/* webpackIgnore: true */ /* turbopackIgnore: true */ 'module');
+      const req = nodeModule.createRequire(path.join(TOOLKIT_ROOT, 'ui', 'package.json'));
+      return req('next/dist/build/swc');
+    })();
+  }
+  return swcPromise;
+};
+
+const transpile = async (file: string, source: string): Promise<string> => {
+  const swc = await getSwc();
+  const ts = /\.tsx?$/.test(file);
+  const out = await swc.transform(source, {
+    filename: file,
+    jsc: {
+      parser: ts ? { syntax: 'typescript', tsx: file.endsWith('x') } : { syntax: 'ecmascript', jsx: true },
+      transform: { react: { runtime: 'automatic', development: false } },
+      target: 'es2020',
+    },
+    module: { type: 'commonjs' },
+    sourceMaps: false,
+    isModule: true,
+  });
+  return out.code as string;
+};
+
+const findUiFile = (pkgDir: string): string | null => {
+  for (const name of UI_FILE_NAMES) {
+    const file = path.join(pkgDir, name);
+    try {
+      if (fs.statSync(file).isFile()) return file;
+    } catch {
+      /* not present */
+    }
+  }
+  return null;
+};
+
+/** All extension UI modules, built-ins first, packages sorted by name. */
+export const listExtensionUiModules = async (): Promise<{
+  modules: ExtensionUiModule[];
+  errors: ExtensionUiError[];
+}> => {
+  const modules: ExtensionUiModule[] = [];
+  const errors: ExtensionUiError[] = [];
+  const seen = new Set<string>();
+  for (const dir of EXTENSION_DIRS) {
+    const root = path.join(TOOLKIT_ROOT, dir);
+    let entries: fs.Dirent[] = [];
+    try {
+      entries = fs.readdirSync(root, { withFileTypes: true });
+    } catch {
+      continue;
+    }
+    const pkgs = entries
+      .filter(e => e.isDirectory() && !e.name.startsWith('.') && !e.name.startsWith('_'))
+      .map(e => e.name)
+      .sort();
+    for (const pkg of pkgs) {
+      const file = findUiFile(path.join(root, pkg));
+      if (!file) continue;
+      const id = `${dir}/${pkg}`;
+      seen.add(id);
+      try {
+        const mtimeMs = fs.statSync(file).mtimeMs;
+        const cached = cache.get(id);
+        if (cached && cached.file === file && cached.mtimeMs === mtimeMs) {
+          modules.push(cached);
+          continue;
+        }
+        const source = await fs.promises.readFile(file, 'utf8');
+        const code = await transpile(file, source);
+        const mod = { id, file, mtimeMs, code };
+        cache.set(id, mod);
+        modules.push(mod);
+      } catch (e: any) {
+        cache.delete(id);
+        errors.push({ id, file, error: e?.message || String(e) });
+      }
+    }
+  }
+  for (const id of cache.keys()) if (!seen.has(id)) cache.delete(id);
+  return { modules, errors };
+};

version.py

@@ -1 +1 @@
-VERSION = "0.13.15"
+VERSION = "0.13.16"