Fix bug

sunho · sunho · commit 1587751dc66b · 2024-08-18T17:36:42.000-07:00
diff --git a/README.md b/README.md
@@ -1,5 +1,5 @@
 # 소개
-Stable Diffusion 계열 API와 모델을 사용하기 편하게 해주는 프론트앤드 프로그램입니다. 모든 씬을 여러번 생성 예약을 해놓고 딴거하다가 와서 이미지를 이상형 월드컵으로 선택하고 리터칭하는 작업 흐름에 맞춰져 있습니다.
+Stable Diffusion 계열 API와 모델을 사용하기 편하게 해주는 프론트앤드 프로그램입니다. 모든 씬을 여러번 생성 예약을 해놓고 딴거하다가 와서 이미지를 이미지 월드컵으로 선택하고 리터칭하는 작업 흐름에 맞춰져 있습니다.
 
 ![](images/img1.png)
 
diff --git a/package.json b/package.json
@@ -19,7 +19,7 @@
     "url": "https://sunho.io"
   },
   "main": "./src/main/main.ts",
-  "version": "3.1.10",
+  "version": "3.1.11",
   "scripts": {
     "build": "concurrently \"npm run build:main\" \"npm run build:renderer\"",
     "build:dll": "cross-env NODE_ENV=development TS_NODE_TRANSPILE_ONLY=true webpack --config ./.erb/configs/webpack.config.renderer.dev.dll.ts",
diff --git a/release/app/package.json b/release/app/package.json
@@ -1,6 +1,6 @@
 {
   "name": "SDStudio",
-  "version": "3.1.10",
+  "version": "3.1.11",
   "description": "A foundation for scalable desktop apps",
   "license": "MIT",
   "author": {
diff --git a/src/main/config.ts b/src/main/config.ts
@@ -21,4 +21,5 @@ export interface Config {
   uuid?: string;
   whiteMode?: boolean;
   disableQuality?: boolean;
+  useAnimalModel?: boolean;
 }
diff --git a/src/renderer/backends/genVendors/nai.ts b/src/renderer/backends/genVendors/nai.ts
@@ -37,26 +37,6 @@ export class NovelAiImageGenService implements ImageGenService {
     return modelMap[model];
   }
 
-  private translateResolution(resolution: Resolution): {
-    height: number;
-    width: number;
-  } {
-    const resolutionMap = {
-      small_landscape: { height: 512, width: 768 },
-      small_portrait: { height: 768, width: 512 },
-      small_square: { height: 640, width: 640 },
-      landscape: { height: 832, width: 1216 },
-      portrait: { height: 1216, width: 832 },
-      square: { height: 1024, width: 1024 },
-      large_landscape: { height: 1024, width: 1536 },
-      large_portrait: { height: 1536, width: 1024 },
-      large_square: { height: 1472, width: 1472 },
-      wallpaper_portrait: { height: 1088, width: 1920 },
-      wallpaper_landscape: { height: 1920, width: 1088 },
-    } as const;
-    return resolutionMap[resolution];
-  }
-
   private translateSampling(sampling: Sampling): string {
     const samplingMap = {
       k_euler_ancestral: 'k_euler_ancestral',
@@ -116,11 +96,16 @@ export class NovelAiImageGenService implements ImageGenService {
   }
 
   public async generateImage(authorization: string, params: ImageGenInput) {
-    const modelValue = this.translateModel(params.model);
-    const resolutionValue = this.translateResolution(params.resolution);
+    let modelValue = this.translateModel(params.model);
+    const resolutionValue = params.resolution;
     const samplingValue = this.translateSampling(params.sampling);
 
     const config = await backend.getConfig();
+    if (config.useAnimalModel) {
+      if (params.model === Model.Anime) {
+        modelValue = 'nai-diffusion-furry-3';
+      }
+    }
 
     const seed = params.seed ?? this.getRandomInt(1, 2100000000);
     let action = undefined;
diff --git a/src/renderer/backends/imageGen.ts b/src/renderer/backends/imageGen.ts
@@ -16,6 +16,7 @@ export enum Resolution {
   LargeSquare = 'large_square',
   WallpaperPortrait = 'wallpaper_portrait',
   WallpaperLandscape = 'wallpaper_landscape',
+  Custom = 'custom',
 }
 
 export const upscaleReoslution = (resolution: Resolution) => {
@@ -36,6 +37,8 @@ export const upscaleReoslution = (resolution: Resolution) => {
       return Resolution.WallpaperPortrait;
     case Resolution.WallpaperLandscape:
       return Resolution.WallpaperLandscape;
+    case Resolution.Custom:
+      return Resolution.Custom;
     default:
       return resolution;
   }
@@ -53,8 +56,16 @@ export const resolutionMap = {
   large_square: { height: 1472, width: 1472 },
   wallpaper_portrait: { height: 1088, width: 1920 },
   wallpaper_landscape: { height: 1920, width: 1088 },
+  custom: { height: 0, width: 0 },
 } as const;
 
+export const convertResolution = (resolution: Resolution): ImageSize => {
+  return {
+    width: resolutionMap[resolution].width,
+    height: resolutionMap[resolution].height,
+  }
+};
+
 export enum Sampling {
   KEulerAncestral = 'k_euler_ancestral',
   KEuler = 'k_euler',
@@ -77,11 +88,16 @@ export interface Vibe {
   strength: number;
 }
 
+export interface ImageSize {
+  width: number;
+  height: number;
+}
+
 export interface ImageGenInput {
   model: Model;
   prompt: string;
   uc: string;
-  resolution: Resolution;
+  resolution: ImageSize;
   sampling: Sampling;
   outputFilePath: string;
   sm: boolean;
diff --git a/src/renderer/componenets/App.css b/src/renderer/componenets/App.css
@@ -345,7 +345,7 @@ input[type='number']::-webkit-outer-spin-button {
 .canvas-tooltip {
   position: absolute;
   right: 0;
-  bottom: 0;
+  bottom: 50px;
   padding: 5px;
   z-index: 3;
   background: white;
diff --git a/src/renderer/componenets/ConfigScreen.tsx b/src/renderer/componenets/ConfigScreen.tsx
@@ -22,6 +22,7 @@ const ConfigScreen = observer(({ onSave }: ConfigScreenProps) => {
   const [whiteMode, setWhiteMode] = useState(false);
   const [noIpCheck, setNoIpCheck] = useState(false);
   const [disableQuality, setDisableQuality] = useState(false);
+  const [animalModel, setAnimalModel] = useState(false);
   const [useLocalBgRemoval, setUseLocalBgRemoval] = useState(false);
   const [refreshImage, setRefreshImage] = useState(false);
   const [ready, setReady] = useState(false);
@@ -41,6 +42,8 @@ const ConfigScreen = observer(({ onSave }: ConfigScreenProps) => {
       setNoIpCheck(config.noIpCheck ?? false);
       setRefreshImage(config.refreshImage ?? false);
       setUseLocalBgRemoval(config.useLocalBgRemoval ?? false);
+      setDisableQuality(config.disableQuality ?? false);
+      setAnimalModel(config.useAnimalModel ?? false);
     })();
     const checkReady = () => {
       setReady(localAIService.ready);
@@ -308,6 +311,16 @@ const ConfigScreen = observer(({ onSave }: ConfigScreenProps) => {
             onChange={(e) => setDisableQuality(e.target.checked)}
           />
         </div>
+        <div className="mt-4 flex items-center gap-2">
+          <label htmlFor="whiteMode" className="text-sm gray-label">
+            NAI 동물 모델 사용
+          </label>
+          <input
+            type="checkbox"
+            checked={animalModel}
+            onChange={(e) => setAnimalModel(e.target.checked)}
+          />
+        </div>
         <button
           className="mt-4 w-full back-sky py-2 rounded hover:brightness-95 active:brightness-90"
           onClick={async () => {
@@ -322,6 +335,7 @@ const ConfigScreen = observer(({ onSave }: ConfigScreenProps) => {
               disableQuality: disableQuality,
               whiteMode: whiteMode,
               useLocalBgRemoval: useLocalBgRemoval,
+              useAnimalModel: animalModel,
             };
             await backend.setConfig(config);
             if (old.useCUDA !== useGPU) localAIService.modelChanged();
diff --git a/src/renderer/componenets/InPaintEditor.tsx b/src/renderer/componenets/InPaintEditor.tsx
@@ -52,6 +52,8 @@ const InPaintEditor = observer(
     const { curSession } = appState;
     const resolutionOptions = Object.entries(resolutionMap)
       .map(([key, value]) => {
+        const resolVal = (editingScene.resolutionWidth ?? '') + 'x' + (editingScene.resolutionHeight ?? '');
+        if (key === 'custom') return { label: '커스텀 (' + resolVal + ')', value: key };
         return { label: `${value.width}x${value.height}`, value: key };
       })
       .filter((x) => !x.value.startsWith('small'));
@@ -138,7 +140,7 @@ const InPaintEditor = observer(
       });
     };
 
-    const confirm = async () => {
+    const saveMask = async () => {
       if (def.hasMask) {
         const mask = await brushTool.current!.getMaskBase64();
         if (!editingScene.preset.mask) {
@@ -154,6 +156,10 @@ const InPaintEditor = observer(
           );
         }
       }
+    };
+
+    const confirm = async () => {
+      await saveMask();
       onConfirm();
     };
 
@@ -191,7 +197,7 @@ const InPaintEditor = observer(
                   options={resolutionOptions}
                   menuPlacement="bottom"
                   selectedOption={editingScene.resolution}
-                  onSelect={(opt) => {
+                  onSelect={async (opt) => {
                     if (
                       opt.value.startsWith('large') ||
                       opt.value.startsWith('wallpaper')
@@ -203,6 +209,25 @@ const InPaintEditor = observer(
                           editingScene.resolution = opt.value as Resolution;
                         },
                       });
+                    } else if (opt.value === 'custom') {
+                      const width = await appState.pushDialogAsync({
+                        type: 'input-confirm',
+                        text: '해상도 너비를 입력해주세요'
+                      });
+                      if (width == null) return;
+                      const height = await appState.pushDialogAsync({
+                        type: 'input-confirm',
+                        text: '해상도 높이를 입력해주세요'
+                      });
+                      if (height == null) return;
+                      try {
+                        const customResolution = { width: parseInt(width), height: parseInt(height) };
+                        editingScene.resolution = opt.value as Resolution;
+                        editingScene.resolutionWidth = (customResolution.width + 63) & ~63;
+                        editingScene.resolutionHeight = (customResolution.height + 63) & ~63;
+                      } catch (e: any) {
+                        appState.pushMessage(e.message);
+                      }
                     } else {
                       editingScene.resolution = opt.value as Resolution;
                     }
@@ -317,6 +342,7 @@ const InPaintEditor = observer(
               <button
                 className={`round-button back-green h-8 w-16 md:w-36 flex items-center justify-center`}
                 onClick={async () => {
+                  await saveMask();
                   await queueI2IWorkflow(curSession!, editingScene.workflowType, editingScene.preset, editingScene, 1, (path: string) => {
                     (async () => {
                       const data = await imageService.fetchImage(path);
diff --git a/src/renderer/componenets/SceneEditor.tsx b/src/renderer/componenets/SceneEditor.tsx
@@ -547,6 +547,8 @@ const SceneEditor = observer(({ scene, onClosed, onDeleted }: Props) => {
 
   const resolutionOptions = Object.entries(resolutionMap)
     .map(([key, value]) => {
+      const resolVal = (scene.resolutionWidth ?? '') + 'x' + (scene.resolutionHeight ?? '');
+      if (key === 'custom') return { label: '커스텀 (' + resolVal + ')', value: key };
       return { label: `${value.width}x${value.height}`, value: key };
     })
     .filter((x) => !x.value.startsWith('small'));
@@ -573,7 +575,7 @@ const SceneEditor = observer(({ scene, onClosed, onDeleted }: Props) => {
                 options={resolutionOptions}
                 menuPlacement="bottom"
                 selectedOption={scene.resolution}
-                onSelect={(opt) => {
+                onSelect={async (opt) => {
                   if (
                     opt.value.startsWith('large') ||
                     opt.value.startsWith('wallpaper')
@@ -585,6 +587,25 @@ const SceneEditor = observer(({ scene, onClosed, onDeleted }: Props) => {
                         scene.resolution = opt.value as Resolution;
                       },
                     });
+                  } else if (opt.value === 'custom') {
+                    const width = await appState.pushDialogAsync({
+                      type: 'input-confirm',
+                      text: '해상도 너비를 입력해주세요'
+                    });
+                    if (width == null) return;
+                    const height = await appState.pushDialogAsync({
+                      type: 'input-confirm',
+                      text: '해상도 높이를 입력해주세요'
+                    });
+                    if (height == null) return;
+                    try {
+                      const customResolution = { width: parseInt(width), height: parseInt(height) };
+                      scene.resolution = opt.value as Resolution;
+                      scene.resolutionWidth = (customResolution.width + 63) & ~63;
+                      scene.resolutionHeight = (customResolution.height + 63) & ~63;
+                    } catch (e: any) {
+                      appState.pushMessage(e.message);
+                    }
                   } else {
                     scene.resolution = opt.value as Resolution;
                   }
diff --git a/src/renderer/componenets/SceneQueueControl.tsx b/src/renderer/componenets/SceneQueueControl.tsx
@@ -596,7 +596,7 @@ const QueueControl = observer(
           const menuItem = oneTimeFlowMap.get(menu)!;
           const input = menuItem.getInput ? await menuItem.getInput(curSession!) : undefined;
           menuItem.handler(curSession!, scene, image, undefined, job, input);
-      }
+        }
       },
     });
 
diff --git a/src/renderer/models/TaskQueueService.ts b/src/renderer/models/TaskQueueService.ts
@@ -1,5 +1,6 @@
 import { v4 } from 'uuid';
 import {
+  convertResolution,
   ImageAugmentInput,
   ImageGenInput,
   Model,
@@ -202,6 +203,17 @@ async function handleNAIDelay(numTry: number, fast: boolean) {
 
 type ImageTaskType = 'gen' | 'inpaint' | 'i2i';
 
+const lowerResolution = (res: Resolution, width?: number, height?: number) => {
+  if (res === Resolution.Custom) {
+    return {
+      width: width!,
+      height: height!,
+    }
+  } else {
+    return convertResolution(res);
+  }
+}
+
 class GenerateImageTaskHandler implements TaskHandler {
   type: ImageTaskType;
   fast: boolean;
@@ -263,11 +275,12 @@ class GenerateImageTaskHandler implements TaskHandler {
         strength: x.strength,
       })),
     );
+    const resol = job.overrideResolution ? job.overrideResolution : task.params.scene!.resolution as Resolution;
     const arg: ImageGenInput = {
       prompt: prompt,
       uc: job.uc,
       model: Model.Anime,
-      resolution: job.overrideResolution ? job.overrideResolution : task.params.scene!.resolution as Resolution,
+      resolution: lowerResolution(resol, task.params.scene!.resolutionWidth, task.params.scene!.resolutionHeight),
       sampling: job.sampling as Sampling,
       sm: job.smea,
       dyn: job.dyn,
@@ -377,6 +390,14 @@ class GenerateImageTaskHandler implements TaskHandler {
         scene: name,
         text: '씬 해상도가 큼',
       });
+    } else if (resolution === Resolution.Custom) {
+      const totalPixels = (task.params.scene.resolutionWidth! * task.params.scene.resolutionHeight!) ?? 0;
+      if (totalPixels > 1024*1024) {
+        res.push({
+          scene: name,
+          text: '씬 해상도가 큼',
+        });
+      }
     }
     return res;
   }
@@ -819,6 +840,7 @@ export const queueI2IWorkflow = async (
   onComplete?: (path: string) => void,
 ) => {
   const def = workFlowService.getDef(type);
+  console.log('queueI2IWorkflow', type, preset, scene, samples, onComplete);
   await def.handler(
     session,
     scene,
diff --git a/src/renderer/models/types.ts b/src/renderer/models/types.ts
diff --git a/src/renderer/models/workflows/SDWorkFlow.ts b/src/renderer/models/workflows/SDWorkFlow.ts

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "SDStudio",`
`3`		`- "version": "3.1.10",`
	`3`	`+ "version": "3.1.11",`
`4`	`4`	`"description": "A foundation for scalable desktop apps",`
`5`	`5`	`"license": "MIT",`
`6`	`6`	`"author": {`
Original file line number	Diff line number	Diff line change
`@@ -21,4 +21,5 @@ export interface Config {`
`21`	`21`	`uuid?: string;`
`22`	`22`	`whiteMode?: boolean;`
`23`	`23`	`disableQuality?: boolean;`
	`24`	`+ useAnimalModel?: boolean;`
`24`	`25`	`}`