Added chroma model to the ui. Added logic to easily pull latest, use local, or use a specific version of chroma. Allow ustom name or path in the ui for custom models

2025-05-07 12:06:30 -06:00
parent d9700bdb99
commit 43cb5603ad
11 changed files with 830 additions and 181 deletions
--- a/ui/src/app/globals.css
+++ b/ui/src/app/globals.css
@@ -19,3 +19,54 @@ body {
  background: var(--background);
  font-family: Arial, Helvetica, sans-serif;
 }
+
+@layer components {
+  /* control */
+  .aitk-react-select-container .aitk-react-select__control {
+    @apply flex w-full h-8 min-h-0 px-0 text-sm bg-gray-800 border border-gray-700 rounded-sm hover:border-gray-600 items-center;
+  }
+
+  /* selected label */
+  .aitk-react-select-container .aitk-react-select__single-value {
+    @apply flex-1 min-w-0 truncate text-sm text-neutral-200;
+  }
+
+  /* invisible input (keeps focus & typing, never wraps) */
+  .aitk-react-select-container .aitk-react-select__input-container {
+    @apply text-neutral-200;
+  }
+
+  /* focus */
+  .aitk-react-select-container .aitk-react-select__control--is-focused {
+    @apply ring-2 ring-gray-600 border-transparent hover:border-transparent shadow-none;
+  }
+
+  /* menu */
+  .aitk-react-select-container .aitk-react-select__menu {
+    @apply bg-gray-800 border border-gray-700;
+  }
+
+  /* options */
+  .aitk-react-select-container .aitk-react-select__option {
+    @apply text-sm text-neutral-200 bg-gray-800 hover:bg-gray-700;
+  }
+
+  /* indicator separator */
+  .aitk-react-select-container .aitk-react-select__indicator-separator {
+    @apply bg-gray-600;
+  }
+
+  /* indicators */
+  .aitk-react-select-container .aitk-react-select__indicators,
+  .aitk-react-select-container .aitk-react-select__indicator {
+    @apply py-0 flex items-center;
+  }
+
+  /* placeholder */
+  .aitk-react-select-container .aitk-react-select__placeholder {
+    @apply text-sm text-neutral-200;
+  }
+}
+
+
+
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/ui/src/app/jobs/new/SimpleJob.tsx
@@ -69,46 +69,6 @@ export default function SimpleJob({

          {/* Model Configuration Section */}
          <Card title="Model Configuration">
-            <SelectInput
-              label="Name or Path"
-              value={jobConfig.config.process[0].model.name_or_path}
-              onChange={value => {
-                // see if model changed
-                const currentModel = options.model.find(
-                  model => model.name_or_path === jobConfig.config.process[0].model.name_or_path,
-                );
-                if (!currentModel || currentModel.name_or_path === value) {
-                  // model has not changed
-                  return;
-                }
-                // revert defaults from previous model
-                for (const key in currentModel.defaults) {
-                  setJobConfig(currentModel.defaults[key][1], key);
-                }
-                // set new model
-                setJobConfig(value, 'config.process[0].model.name_or_path');
-                // update the defaults when a model is selected
-                const model = options.model.find(model => model.name_or_path === value);
-                if (model?.defaults) {
-                  for (const key in model.defaults) {
-                    setJobConfig(model.defaults[key][0], key);
-                  }
-                }
-              }}
-              options={
-                options.model
-                  .map(model => {
-                    if (model.dev_only && !isDev) {
-                      return null;
-                    }
-                    return {
-                      value: model.name_or_path,
-                      label: model.name_or_path,
-                    };
-                  })
-                  .filter(x => x) as { value: string; label: string }[]
-              }
-            />
            <SelectInput
              label="Model Architecture"
              value={jobConfig.config.process[0].model.arch}
@@ -117,6 +77,18 @@ export default function SimpleJob({
                if (!currentArch || currentArch.name === value) {
                  return;
                }
+
+                // revert defaults from previous model
+                for (const key in currentArch.defaults) {
+                  setJobConfig(currentArch.defaults[key][1], key);
+                }
+                // update the defaults when a model is selected
+                const newArch = modelArchs.find(model => model.name === value);
+                if (newArch?.defaults) {
+                  for (const key in newArch.defaults) {
+                    setJobConfig(newArch.defaults[key][0], key);
+                  }
+                }
                // set new model
                setJobConfig(value, 'config.process[0].model.arch');
              }}
@@ -131,6 +103,18 @@ export default function SimpleJob({
                  .filter(x => x) as { value: string; label: string }[]
              }
            />
+            <TextInput
+              label="Name or Path"
+              value={jobConfig.config.process[0].model.name_or_path}
+              onChange={(value: string | null) => {
+                if (value?.trim() === '') {
+                  value = null;
+                }
+                setJobConfig(value, 'config.process[0].model.name_or_path');
+              }}
+              placeholder=""
+              required
+            />
            <FormGroup label="Quantize">
              <div className="grid grid-cols-2 gap-2">
                <Checkbox
--- a/ui/src/app/jobs/new/jobConfig.ts
+++ b/ui/src/app/jobs/new/jobConfig.ts
@@ -73,7 +73,7 @@ export const defaultJobConfig: JobConfig = {
          name_or_path: 'ostris/Flex.1-alpha',
          quantize: true,
          quantize_te: true,
-          arch: 'flux',
+          arch: 'flex1',
          low_vram: false,
        },
        sample: {
--- a/ui/src/app/jobs/new/options.ts
+++ b/ui/src/app/jobs/new/options.ts
@@ -1,5 +1,6 @@
 export interface Model {
  name_or_path: string;
+  arch: string;
  dev_only?: boolean;
  defaults?: { [key: string]: any };
 }
@@ -8,11 +9,108 @@ export interface Option {
  model: Model[];
 }

-export const modelArchs = [
-  { name: 'flux', label: 'Flux.1' },
-  { name: 'wan21', label: 'Wan 2.1' },
-  { name: 'lumina2', label: 'Lumina2' },
-  { name: 'hidream', label: 'HiDream' },
+export interface ModelArch {
+  name: string;
+  label: string;
+  defaults?: { [key: string]: [any, any] };
+}
+
+const defaultNameOrPath = '';
+
+export const modelArchs: ModelArch[] = [
+  {
+    name: 'flux',
+    label: 'FLUX.1',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['black-forest-labs/FLUX.1-dev', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+  },
+  {
+    name: 'flex1',
+    label: 'Flex.1',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['ostris/Flex.1-alpha', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].train.bypass_guidance_embedding': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+  },
+  // { name: 'flex2', label: 'Flex.2' },
+  {
+    name: 'chroma',
+    label: 'Chroma',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['lodestones/Chroma', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+  },
+  {
+    name: 'wan21:1b',
+    label: 'Wan 2.1 (1.3B)',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Wan-AI/Wan2.1-T2V-1.3B-Diffusers', defaultNameOrPath],
+      'config.process[0].model.quantize': [false, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [40, 1],
+      'config.process[0].sample.fps': [15, 1],
+    },
+  },
+  {
+    name: 'wan21:14b',
+    label: 'Wan 2.1 (14B)',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Wan-AI/Wan2.1-T2V-14B-Diffuserss', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [40, 1],
+      'config.process[0].sample.fps': [15, 1],
+    },
+  },
+  {
+    name: 'lumina2',
+    label: 'Lumina2',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Alpha-VLLM/Lumina-Image-2.0', defaultNameOrPath],
+      'config.process[0].model.quantize': [false, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+    },
+  },
+  {
+    name: 'hidream',
+    label: 'HiDream',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['HiDream-ai/HiDream-I1-Full', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.lr': [0.0002, 0.0001],
+      'config.process[0].train.timestep_type': ['shift', 'sigmoid'],
+      'config.process[0].network.network_kwargs.ignore_if_contains': [['ff_i.experts', 'ff_i.gate'], []],
+    },
+  },
 ];

 export const isVideoModelFromArch = (arch: string) => {
@@ -22,15 +120,15 @@ export const isVideoModelFromArch = (arch: string) => {

 const defaultModelArch = 'flux';

-export const options = {
+export const options: Option = {
  model: [
    {
      name_or_path: 'ostris/Flex.1-alpha',
+      arch: 'flex1',
      defaults: {
        // default updates when [selected, unselected] in the UI
        'config.process[0].model.quantize': [true, false],
        'config.process[0].model.quantize_te': [true, false],
-        'config.process[0].model.arch': ['flux', defaultModelArch],
        'config.process[0].train.bypass_guidance_embedding': [true, false],
        'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
        'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
@@ -38,22 +136,33 @@ export const options = {
    },
    {
      name_or_path: 'black-forest-labs/FLUX.1-dev',
+      arch: 'flux',
+      defaults: {
+        // default updates when [selected, unselected] in the UI
+        'config.process[0].model.quantize': [true, false],
+        'config.process[0].model.quantize_te': [true, false],
+        'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+        'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      },
+    },
+    {
+      name_or_path: 'lodestones/Chroma',
+      arch: 'chroma',
      defaults: {
        // default updates when [selected, unselected] in the UI
        'config.process[0].model.quantize': [true, false],
        'config.process[0].model.quantize_te': [true, false],
-        'config.process[0].model.arch': ['flux', defaultModelArch],
        'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
        'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
      },
    },
    {
      name_or_path: 'Wan-AI/Wan2.1-T2V-1.3B-Diffusers',
+      arch: 'wan21',
      defaults: {
        // default updates when [selected, unselected] in the UI
        'config.process[0].model.quantize': [false, false],
        'config.process[0].model.quantize_te': [true, false],
-        'config.process[0].model.arch': ['wan21', defaultModelArch],
        'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
        'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
        'config.process[0].sample.num_frames': [40, 1],
@@ -62,11 +171,11 @@ export const options = {
    },
    {
      name_or_path: 'Wan-AI/Wan2.1-T2V-14B-Diffusers',
+      arch: 'wan21',
      defaults: {
        // default updates when [selected, unselected] in the UI
        'config.process[0].model.quantize': [true, false],
        'config.process[0].model.quantize_te': [true, false],
-        'config.process[0].model.arch': ['wan21', defaultModelArch],
        'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
        'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
        'config.process[0].sample.num_frames': [40, 1],
@@ -75,22 +184,22 @@ export const options = {
    },
    {
      name_or_path: 'Alpha-VLLM/Lumina-Image-2.0',
+      arch: 'lumina2',
      defaults: {
        // default updates when [selected, unselected] in the UI
        'config.process[0].model.quantize': [false, false],
        'config.process[0].model.quantize_te': [true, false],
-        'config.process[0].model.arch': ['lumina2', defaultModelArch],
        'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
        'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
      },
    },
    {
      name_or_path: 'HiDream-ai/HiDream-I1-Full',
+      arch: 'hidream',
      defaults: {
        // default updates when [selected, unselected] in the UI
        'config.process[0].model.quantize': [true, false],
        'config.process[0].model.quantize_te': [true, false],
-        'config.process[0].model.arch': ['hidream', defaultModelArch],
        'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
        'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
        'config.process[0].train.lr': [0.0002, 0.0001],
@@ -100,11 +209,11 @@ export const options = {
    },
    {
      name_or_path: 'ostris/objective-reality',
+      arch: 'sd1',
      dev_only: true,
      defaults: {
        'config.process[0].sample.sampler': ['ddpm', 'flowmatch'],
        'config.process[0].train.noise_scheduler': ['ddpm', 'flowmatch'],
-        'config.process[0].model.arch': ['sd1', defaultModelArch],
      },
    },
  ],
--- a/ui/src/components/formInputs.tsx
+++ b/ui/src/components/formInputs.tsx
@@ -1,5 +1,9 @@
+'use client';
+
 import React, { forwardRef } from 'react';
 import classNames from 'classnames';
+import dynamic from "next/dynamic";
+const Select = dynamic(() => import("react-select"), { ssr: false });

 const labelClasses = 'block text-xs mb-1 mt-2 text-gray-300';
 const inputClasses =
@@ -116,16 +120,21 @@ export interface SelectInputProps extends InputProps {

 export const SelectInput = (props: SelectInputProps) => {
  const { label, value, onChange, options } = props;
+  const selectedOption = options.find(option => option.value === value);
  return (
    <div className={classNames(props.className)}>
      {label && <label className={labelClasses}>{label}</label>}
-      <select value={value} onChange={e => onChange(e.target.value)} className={inputClasses}>
-        {options.map(option => (
-          <option key={option.value} value={option.value}>
-            {option.label}
-          </option>
-        ))}
-      </select>
+      <Select 
+        value={selectedOption} 
+        options={options}
+        className="aitk-react-select-container"
+        classNamePrefix="aitk-react-select"
+        onChange={selected => {
+          if (selected) {
+            onChange((selected as { value: string }).value);
+          }
+        }}
+      />
    </div>
  );
 };