From d5cda1ba4236dc26485ef99896978009da30fa70 Mon Sep 17 00:00:00 2001
From: Dave Schumaker <daveschumaker@users.noreply.github.com>
Date: Sun, 1 Sep 2024 09:11:48 -0700
Subject: [PATCH] fix: Update samplers to use type and rename
 'k_euler_ancestral' to 'k_euler_a'. (#50)

---
 .../AdvancedOptions/SamplerSelect.tsx         | 18 +----
 .../ImageParamsForHordeApi.test.ts            | 70 +++++++++----------
 app/_data-models/ImageParamsForHordeApi.ts    |  3 +-
 app/_data-models/PromptInput.test.ts          |  2 +-
 app/_data-models/PromptInput.ts               |  3 +-
 app/_types/HordeTypes.ts                      | 15 ++++
 6 files changed, 57 insertions(+), 54 deletions(-)

diff --git a/app/_components/AdvancedOptions/SamplerSelect.tsx b/app/_components/AdvancedOptions/SamplerSelect.tsx
index 1abd34b..6a3c15b 100644
--- a/app/_components/AdvancedOptions/SamplerSelect.tsx
+++ b/app/_components/AdvancedOptions/SamplerSelect.tsx
@@ -3,21 +3,7 @@
 import { useInput } from '@/app/_providers/PromptInputProvider'
 import Select, { SelectOption } from '../Select'
 import OptionLabel from './OptionLabel'
-
-export type SamplerOption =
-  | 'DDIM'
-  | 'k_dpm_2_a'
-  | 'k_dpm_2'
-  | 'k_dpm_adaptive'
-  | 'k_dpm_fast'
-  | 'k_dpmpp_2m'
-  | 'k_dpmpp_2s_a'
-  | 'k_dpmpp_sde'
-  | 'k_euler_a'
-  | 'k_euler'
-  | 'k_heun'
-  | 'k_lms'
-  | 'lcm'
+import { SamplerOption } from '@/app/_types/HordeTypes'
 
 const samplers: Array<{ value: SamplerOption; label: SamplerOption }> = [
   { value: 'DDIM', label: 'DDIM' },
@@ -47,7 +33,7 @@ export default function SamplerSelect() {
       <div className="w-full">
         <Select
           onChange={(option: SelectOption) => {
-            setInput({ sampler: option.value as string })
+            setInput({ sampler: option.value as SamplerOption })
           }}
           options={samplers.map((sampler) => ({
             value: sampler.value,
diff --git a/app/_data-models/ImageParamsForHordeApi.test.ts b/app/_data-models/ImageParamsForHordeApi.test.ts
index 890ca81..051b16d 100644
--- a/app/_data-models/ImageParamsForHordeApi.test.ts
+++ b/app/_data-models/ImageParamsForHordeApi.test.ts
@@ -30,40 +30,40 @@ describe('ImageParamsForHordeApi', () => {
       height: 512,
       width: 512,
       steps: 30,
-      sampler: 'k_euler_ancestral',
+      sampler: 'k_euler_a',
       seed: '123456',
       models: ['stable_diffusion']
     })
 
-    // Mock AppSettings.get
-    ;(AppSettings.get as jest.Mock).mockImplementation(
-      (key: AppSettingsKey) => {
-        const settings: AppSettingsParams = {
-          allowedWorkers: [],
-          allowNsfwImages: false,
-          apiKey: 'test-api-key',
-          autoDowngrade: false,
-          blockedWorkers: [],
-          civitAiBaseModelFilter: [],
-          negativePanelOpen: false,
-          runInBackground: false,
-          saveInputOnCreate: false,
-          sharedKey: '',
-          slow_workers: true,
-          useAllowedWorkers: false,
-          useBeta: false,
-          useBlockedWorkers: false,
-          useReplacementFilter: true,
-          useTrusted: true
+      // Mock AppSettings.get
+      ; (AppSettings.get as jest.Mock).mockImplementation(
+        (key: AppSettingsKey) => {
+          const settings: AppSettingsParams = {
+            allowedWorkers: [],
+            allowNsfwImages: false,
+            apiKey: 'test-api-key',
+            autoDowngrade: false,
+            blockedWorkers: [],
+            civitAiBaseModelFilter: [],
+            negativePanelOpen: false,
+            runInBackground: false,
+            saveInputOnCreate: false,
+            sharedKey: '',
+            slow_workers: true,
+            useAllowedWorkers: false,
+            useBeta: false,
+            useBlockedWorkers: false,
+            useReplacementFilter: true,
+            useTrusted: true
+          }
+
+          if (key in settings) {
+            return settings[key]
+          }
+
+          throw new Error(`Unexpected key: ${key}`)
         }
-
-        if (key in settings) {
-          return settings[key]
-        }
-
-        throw new Error(`Unexpected key: ${key}`)
-      }
-    )
+      )
   })
 
   test('setBaseParams should set correct base parameters', () => {
@@ -77,7 +77,7 @@ describe('ImageParamsForHordeApi', () => {
         height: 512,
         width: 512,
         steps: 30,
-        sampler_name: 'k_euler_ancestral',
+        sampler_name: 'k_euler_a',
         seed: '123456'
       },
       nsfw: false,
@@ -141,12 +141,12 @@ describe('ImageParamsForHordeApi', () => {
 
   test('setSourceProcessing should set correct source processing parameters', async () => {
     // Mock the getImagesForArtbotJobFromDexie function
-    ;(dbModule.getImagesForArtbotJobFromDexie as jest.Mock).mockResolvedValue([
+    ; (dbModule.getImagesForArtbotJobFromDexie as jest.Mock).mockResolvedValue([
       { imageBlobBuffer: new ArrayBuffer(8) }
     ])
 
-    // Mock blobToBase64
-    ;(imageUtils.blobToBase64 as jest.Mock).mockResolvedValue('base64string')
+      // Mock blobToBase64
+      ; (imageUtils.blobToBase64 as jest.Mock).mockResolvedValue('base64string')
 
     promptInput.source_processing = SourceProcessing.Img2Img
     promptInput.denoising_strength = 0.6
@@ -178,7 +178,7 @@ describe('ImageParamsForHordeApi', () => {
         height: 512,
         width: 512,
         steps: 30,
-        sampler_name: 'k_euler_ancestral',
+        sampler_name: 'k_euler_a',
         seed: '123456',
         n: 1,
         post_processing: [],
@@ -205,7 +205,7 @@ describe('ImageParamsForHordeApi', () => {
     expect(result.height).toBe(512)
     expect(result.width).toBe(512)
     expect(result.steps).toBe(30)
-    expect(result.sampler).toBe('k_euler_ancestral')
+    expect(result.sampler).toBe('k_euler_a')
     expect(result.seed).toBe('123456')
     expect(result.models).toEqual(['stable_diffusion'])
   })
diff --git a/app/_data-models/ImageParamsForHordeApi.ts b/app/_data-models/ImageParamsForHordeApi.ts
index 78bdcc5..4e55092 100644
--- a/app/_data-models/ImageParamsForHordeApi.ts
+++ b/app/_data-models/ImageParamsForHordeApi.ts
@@ -3,6 +3,7 @@ import {
   ControlTypes,
   HordeTi,
   Lora,
+  SamplerOption,
   SourceProcessing
 } from '../_types/HordeTypes'
 import { castTiInject } from '../_utils/hordeUtils'
@@ -26,7 +27,7 @@ interface HordeApiParamsBuilderInterface {
 }
 
 export interface ImageParams {
-  sampler_name?: string // Optional due to ControlNet
+  sampler_name?: SamplerOption // Optional due to ControlNet
   cfg_scale: number
   height: number
   width: number
diff --git a/app/_data-models/PromptInput.test.ts b/app/_data-models/PromptInput.test.ts
index fea10ff..ed80d11 100644
--- a/app/_data-models/PromptInput.test.ts
+++ b/app/_data-models/PromptInput.test.ts
@@ -31,7 +31,7 @@ describe('PromptInput', () => {
     expect(defaultPromptInput.preset).toEqual([])
     expect(defaultPromptInput.prompt).toBe('')
     expect(defaultPromptInput.return_control_map).toBe(false)
-    expect(defaultPromptInput.sampler).toBe('euler_a')
+    expect(defaultPromptInput.sampler).toBe('k_euler_a')
     expect(defaultPromptInput.seed).toBe('')
     expect(defaultPromptInput.source_processing).toBe(SourceProcessing.Prompt)
     expect(defaultPromptInput.steps).toBe(8)
diff --git a/app/_data-models/PromptInput.ts b/app/_data-models/PromptInput.ts
index 108d3d8..54f2ed5 100644
--- a/app/_data-models/PromptInput.ts
+++ b/app/_data-models/PromptInput.ts
@@ -1,6 +1,7 @@
 import { ImageOrientations, JobType, Workflow } from '@/app/_types/ArtbotTypes'
 import {
   ControlTypes,
+  SamplerOption,
   SourceProcessing,
   StylePresetConfig
 } from '@/app/_types/HordeTypes'
@@ -64,7 +65,7 @@ class PromptInput {
   post_processing: Array<string> = []
   prompt: string = ''
   return_control_map: boolean = false
-  sampler: string = 'euler_a'
+  sampler: SamplerOption = 'k_euler_a'
   seed: string = ''
   source_processing?: SourceProcessing = SourceProcessing.Prompt
   steps: number = 8
diff --git a/app/_types/HordeTypes.ts b/app/_types/HordeTypes.ts
index 1de87af..bce5d18 100644
--- a/app/_types/HordeTypes.ts
+++ b/app/_types/HordeTypes.ts
@@ -239,6 +239,21 @@ export interface LoraConfig {
   is_version: boolean
 }
 
+export type SamplerOption =
+  | 'DDIM'
+  | 'k_dpm_2_a'
+  | 'k_dpm_2'
+  | 'k_dpm_adaptive'
+  | 'k_dpm_fast'
+  | 'k_dpmpp_2m'
+  | 'k_dpmpp_2s_a'
+  | 'k_dpmpp_sde'
+  | 'k_euler_a'
+  | 'k_euler'
+  | 'k_heun'
+  | 'k_lms'
+  | 'lcm'
+
 export interface SharedApiKey {
   id: string
   kudos: number