diff --git a/extensions_built_in/audio_models/ui.tsx b/extensions_built_in/audio_models/ui.tsx
new file mode 100644
index 0000000..a86eec9
--- /dev/null
+++ b/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: (
+
+
+ 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.
+
+
+ YuE2 is two experts on one backbone. The AR expert reads the
+ style line and lyrics and writes the song as a sequence of semantic
+ codec tokens (25 per second). The NAR expert 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.
+
+
+ The AR loss is the one to watch (loss/ar_ce). It starts
+ near 5 and, on a small dataset, falls toward 0 within a few hundred
+ steps, which is memorization. ar_kl_weight in model
+ kwargs anchors the AR to the base model so it cannot collapse onto the
+ training songs; ar_lr_multiplier 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.
+
+
+ The official audio-to-token encoder is unreleased. Training uses the
+ community tokenizer by Kytra (
+
+ @sin_ceriously
+
+ ), a MERT-v2-FullSong head that maps real audio to YuE2 codec tokens:{" "}
+
+ Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4
+
+ . 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.
+
+
+ ABC generation. 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{" "}
+ cot: full (chords) or cot: melody (melody
+ only), or the single stage with cot: off.
+
+
+ Training for both. Every song is transcribed to an ABC sheet
+ with SheetSage2 and the AR is trained on lyrics to sheet and sheet to
+ tokens. abc_dropout (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.
+
+
+ Caching is required. 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 cot means
+ deleting the dataset's _latent_cache folder so the sheets
+ are rebuilt. Audio has no resolution, so keep a single resolution
+ bucket per dataset or every bucket duplicates the songs.
+
+
+ Prompt format: a style line, then [Lyrics], then the
+ lyrics with bracketed section headers such as [Verse 1]{" "}
+ and [Chorus]. 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.
+
+
+ MiniMax H3 is a guidance distilled model, so training on it
+ directly will make the guidance distillation break down. There are
+ two different ways to train on this model without breaking the
+ guidance distillation: Contrastive Guidance and Training Adapter.
+ Both have their pros and cons. The adapter is faster, but will
+ still break down over a long run. Contrastive Guidance is slower,
+ but is less likely to break down.
+
+ ),
+ },
+ },
+ ],
+ modelNotes: (
+
+
+ Weights load from the{" "}
+ Comfy-Org/MiniMax-H3 on first load (~43GB total). Files
+ used:
+
+
+ {`/
+├── diffusion_models/
+│ └── minimax_h3_fl2va_pruned_int8_convrot.safetensors
+├── text_encoders/
+│ └── qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors
+└── vae/
+ ├── minimax_h3_video_vae_fp16.safetensors
+ └── minimax_h3_audio_vae_fp32.safetensors`}
+
+
+ The checkpoints are pre-quantized and load directly: int8 ConvRot DiT
+ (~21GB) and nvfp4 Qwen3-VL text encoder (~16GB). The default qtypes (
+ convrot8 / nvfp4) match the files exactly,
+ so nothing is re-quantized on load. Picking a different quantization
+ re-quantizes the pre-quantized layers into that format, one layer at a
+ time.
+
+
+ Supports t2v and first-frame i2v (ctrl img / i2v datasets) with joint
+ audio. The model is guidance-distilled — keep guidance scale at 1.
+ Video is fixed 24 fps and frame counts snap down to the 17n+5 grid (5,
+ 22, 39, 56, ..., 107, 124 ≈ 5s). Image datasets (num_frames 1) train
+ as single frames, and a sample with num_frames 1 renders a single
+ image.
+
+
+ MiniMax H3 is a guidance distilled model, so training on it
+ directly will make the guidance distillation break down. There are
+ two different ways to train on this model without breaking the
+ guidance distillation: Contrastive Guidance and Training Adapter.
+ Both have their pros and cons. The adapter is faster, but will
+ still break down over a long run. Contrastive Guidance is slower,
+ but is less likely to break down. D-OPSD instead self-distills: a
+ no-grad teacher pass sees the training target as its own reference
+ and its prediction becomes the training target for a
+ reference-free pass, baking the reference into your trigger word
+ (or the caption itself when no trigger word is set). Re-caches
+ latents with pixel tensors.
+
+ ),
+ },
+ },
+ {
+ label: "Image Reference Presentation",
+ options: [
+ { value: "picture", label: "Picture (default)" },
+ { value: "video", label: "Static video clip" },
+ ],
+ getValue: (config: JobConfig) => {
+ return config?.config?.process?.[0]?.model?.model_kwargs
+ ?.image_refs_as_video
+ ? "video"
+ : "picture";
+ },
+ onChange: (
+ value: string,
+ config: JobConfig,
+ setJobConfig: (value: any, key: string) => void,
+ ) => {
+ const kwargs = {
+ ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}),
+ };
+ if (value === "video") {
+ kwargs.image_refs_as_video = true;
+ } else {
+ delete kwargs.image_refs_as_video;
+ delete kwargs.image_ref_video_frames;
+ }
+ setJobConfig(kwargs, "config.process[0].model.model_kwargs");
+ },
+ doc: {
+ title: "MiniMax-H3 Image Reference Presentation",
+ description: (
+
+
+ How still-image references (dataset control images and sample
+ ctrl images) are shown to the model. Video references always use
+ the video path.
+
+
+ Picture : the native ref2va recipe — a
+ single-frame reference block, shown to Qwen3-VL as a{" "}
+ <Picture i> block, scaled down only.
+
+
+ Static video clip : the image is held for 5
+ frames (2 latent frames) and routed through the exact path a
+ reference VIDEO takes — video sizing (matched to the target's
+ pixel area), multi-frame reference block,{" "}
+ <Video k> timestamped presentation. Use this
+ when training on image references but sampling with video
+ references, so the LoRA learns the pathway it will be used
+ through. Adds a handful of rows per reference. Frame count is
+ adjustable with model_kwargs.image_ref_video_frames{" "}
+ (17n+5). Changing this re-caches text embeddings.
+
+
+
+ Reference-to-video: control images and videos condition the output as
+ subject/style references (never as a first frame). References keep
+ their own aspect and are matched to the target's pixel area (images
+ scale down only, never up; a same-aspect video reference is exactly
+ the target size). Each rides into the packed sequence as a reference
+ block, and is also shown to the Qwen3-VL conditioner as a{" "}
+ <Picture i> (image) or timestamped{" "}
+ <Video k> (video) vision block. Training references
+ come from the dataset control path(s); sampling uses the sample ctrl
+ images — always as references. The Image Reference Presentation option
+ can route still images through the video-reference path as short
+ static clips.
+
+
+ Weights load like MiniMax-H3 (see that arch's notes) from the{" "}
+
+ diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors
+ {" "}
+ — the ref2va partition of the same release; text encoder and VAEs are
+ shared with the fl2va arch. Everything else (pre-quantized load, 24
+ fps, 17n+5 frame grid, guidance scale 1, single-image mode) matches
+ MiniMax-H3.
+
+
+
+ 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).
+
+
+ The instruction in model_kwargs.instruction 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.
+
+
+ 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 loss/ce: next-token cross-entropy on
+ the caption.
+
+
- MiniMax H3 is a guidance distilled model, so training on it directly will make the guidance distillation
- break down. There are two different ways to train on this model without breaking the guidance
- distillation: Contrastive Guidance and Training Adapter. Both have their pros and cons. The adapter is
- faster, but will still break down over a long run. Contrastive Guidance is slower, but is less likely to
- break down.
-
- ),
- },
- },
- ],
- modelNotes: (
-
-
- Weights load from the{' '}
- Comfy-Org/MiniMax-H3 on first load
- (~43GB total). Files used:
-
-
- {`/
-├── diffusion_models/
-│ └── minimax_h3_fl2va_pruned_int8_convrot.safetensors
-├── text_encoders/
-│ └── qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors
-└── vae/
- ├── minimax_h3_video_vae_fp16.safetensors
- └── minimax_h3_audio_vae_fp32.safetensors`}
-
-
- The checkpoints are pre-quantized and load directly: int8 ConvRot DiT (~21GB) and nvfp4 Qwen3-VL text encoder
- (~16GB). The default qtypes (convrot8 / nvfp4) match the files exactly, so nothing
- is re-quantized on load. Picking a different quantization re-quantizes the pre-quantized layers into that
- format, one layer at a time.
-
-
- Supports t2v and first-frame i2v (ctrl img / i2v datasets) with joint audio. The model is guidance-distilled —
- keep guidance scale at 1. Video is fixed 24 fps and frame counts snap down to the 17n+5 grid (5, 22, 39, 56,
- ..., 107, 124 ≈ 5s). Image datasets (num_frames 1) train as single frames, and a sample with num_frames 1
- renders a single image.
-
-
- MiniMax H3 is a guidance distilled model, so training on it directly will make the guidance distillation
- break down. There are two different ways to train on this model without breaking the guidance
- distillation: Contrastive Guidance and Training Adapter. Both have their pros and cons. The adapter is
- faster, but will still break down over a long run. Contrastive Guidance is slower, but is less likely to
- break down. D-OPSD instead self-distills: a no-grad teacher pass sees the training target as its own
- reference and its prediction becomes the training target for a reference-free pass, baking the reference
- into your trigger word (or the caption itself when no trigger word is set). Re-caches latents with pixel
- tensors.
-
- ),
- },
- },
- {
- label: 'Image Reference Presentation',
- options: [
- { value: 'picture', label: 'Picture (default)' },
- { value: 'video', label: 'Static video clip' },
- ],
- getValue: (config: JobConfig) => {
- return config?.config?.process?.[0]?.model?.model_kwargs?.image_refs_as_video ? 'video' : 'picture';
- },
- onChange: (value: string, config: JobConfig, setJobConfig: (value: any, key: string) => void) => {
- const kwargs = { ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}) };
- if (value === 'video') {
- kwargs.image_refs_as_video = true;
- } else {
- delete kwargs.image_refs_as_video;
- delete kwargs.image_ref_video_frames;
- }
- setJobConfig(kwargs, 'config.process[0].model.model_kwargs');
- },
- doc: {
- title: 'MiniMax-H3 Image Reference Presentation',
- description: (
-
-
- How still-image references (dataset control images and sample ctrl images) are shown to the model. Video
- references always use the video path.
-
-
- Picture : the native ref2va recipe — a single-frame reference block, shown to Qwen3-VL
- as a <Picture i> block, scaled down only.
-
-
- Static video clip : the image is held for 5 frames (2 latent frames) and routed through
- the exact path a reference VIDEO takes — video sizing (matched to the target's pixel area), multi-frame
- reference block, <Video k> timestamped presentation. Use this when training on image
- references but sampling with video references, so the LoRA learns the pathway it will be used through.
- Adds a handful of rows per reference. Frame count is adjustable with{' '}
- model_kwargs.image_ref_video_frames (17n+5). Changing this re-caches text embeddings.
-
-
-
- Reference-to-video: control images and videos condition the output as subject/style references (never as a
- first frame). References keep their own aspect and are matched to the target's pixel area (images scale down
- only, never up; a same-aspect video reference is exactly the target size). Each rides into the packed sequence
- as a reference block, and is also shown to the Qwen3-VL conditioner as a <Picture i>{' '}
- (image) or timestamped <Video k> (video) vision block. Training references come from the
- dataset control path(s); sampling uses the sample ctrl images — always as references. The Image Reference
- Presentation option can route still images through the video-reference path as short static clips.
-
-
- Weights load like MiniMax-H3 (see that arch's notes) from the{' '}
- diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors — the ref2va partition
- of the same release; text encoder and VAEs are shared with the fl2va arch. Everything else (pre-quantized
- load, 24 fps, 17n+5 frame grid, guidance scale 1, single-image mode) matches MiniMax-H3.
-
-
-
- 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).
-
-
- The instruction in model_kwargs.instruction 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.
-
-
- 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 loss/ce: next-token cross-entropy on the caption.
-
-
-
- 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.
-
-
- YuE2 is two experts on one backbone. The AR expert reads the style line and lyrics and writes the song
- as a sequence of semantic codec tokens (25 per second). The NAR expert 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.
-
-
- The AR loss is the one to watch (loss/ar_ce). It starts near 5 and, on a small dataset, falls
- toward 0 within a few hundred steps, which is memorization. ar_kl_weight in model kwargs anchors
- the AR to the base model so it cannot collapse onto the training songs; ar_lr_multiplier 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.
-
-
- The official audio-to-token encoder is unreleased. Training uses the community tokenizer by Kytra (
-
- @sin_ceriously
-
- ), a MERT-v2-FullSong head that maps real audio to YuE2 codec tokens:{' '}
-
- Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4
-
- . 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.
-
-
- ABC generation. 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 cot: full (chords) or cot: melody (melody only), or the single
- stage with cot: off.
-
-
- Training for both. Every song is transcribed to an ABC sheet with SheetSage2 and the AR is trained on
- lyrics to sheet and sheet to tokens. abc_dropout (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.
-
-
- Caching is required. 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{' '}
- cot means deleting the dataset's _latent_cache folder so the sheets are rebuilt.
- Audio has no resolution, so keep a single resolution bucket per dataset or every bucket duplicates the songs.
-
-
- Prompt format: a style line, then [Lyrics], then the lyrics with bracketed section headers such
- as [Verse 1] and [Chorus]. 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.
-
-
/__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//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).
diff --git a/ui/src/extensions/modelArchs.ts b/ui/src/extensions/modelArchs.ts
new file mode 100644
index 0000000..af97cd4
--- /dev/null
+++ b/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 ''` 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({ archs: [], errors: [], isLoaded: false });
+let inflight: Promise | 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 => {
+ 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();
+ 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;
+} => {
+ const [current] = state.use();
+ useEffect(() => {
+ loadModelArchs();
+ }, []);
+ const groupedModelOptions = useMemo(() => options.groupModelOptions(current.archs), [current.archs]);
+ return { ...current, groupedModelOptions, refresh: loadModelArchs };
+};
diff --git a/ui/src/server/extensionUi.ts b/ui/src/server/extensionUi.ts
new file mode 100644
index 0000000..6ee0366
--- /dev/null
+++ b/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; // "/"
+ file: string; // absolute path
+ mtimeMs: number;
+ code: string; // CommonJS source
+}
+
+export interface ExtensionUiError {
+ id: string;
+ file: string;
+ error: string;
+}
+
+const cache = new Map();
+
+// 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 | 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 => {
+ 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();
+ 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 };
+};
diff --git a/version.py b/version.py
index 0156c0c..5c22c0c 100644
--- a/version.py
+++ b/version.py
@@ -1 +1 @@
-VERSION = "0.13.15"
+VERSION = "0.13.16"