Adds Codeformer support

2026-04-23 03:00:31 -04:00 · 2022-10-18 21:57:40 +08:00
parent 62dd3b7d7d
commit 32a5ffe436
15 changed files with 915 additions and 675 deletions
--- a/frontend/src/app/constants.ts
+++ b/frontend/src/app/constants.ts
@@ -43,7 +43,7 @@ export const PARAMETERS: { [key: string]: string } = {
  sampler: 'Sampler',
  seed: 'Seed',
  img2imgStrength: 'img2img Strength',
-  gfpganStrength: 'GFPGAN Strength',
+  facetoolStrength: 'GFPGAN Strength',
  upscalingLevel: 'Upscaling Level',
  upscalingStrength: 'Upscaling Strength',
  initialImagePath: 'Initial Image',
@@ -56,3 +56,5 @@ export const PARAMETERS: { [key: string]: string } = {
 export const NUMPY_RAND_MIN = 0;

 export const NUMPY_RAND_MAX = 4294967295;
+
+export const FACETOOL_TYPES = ['gfpgan', 'codeformer'] as const;
--- a/frontend/src/app/invokeai.d.ts
+++ b/frontend/src/app/invokeai.d.ts
@@ -89,15 +89,16 @@ export declare type ESRGANMetadata = CommonPostProcessedImageMetadata & {
  strength: number;
 };

-export declare type GFPGANMetadata = CommonPostProcessedImageMetadata & {
-  type: 'gfpgan';
+export declare type FacetoolMetadata = CommonPostProcessedImageMetadata & {
+  type: 'gfpgan' | 'codeformer';
  strength: number;
+  fidelity?: number;
 };

 // Superset of all postprocessed image metadata types..
 export declare type PostProcessedImageMetadata =
  | ESRGANMetadata
-  | GFPGANMetadata;
+  | FacetoolMetadata;

 // Metadata includes the system config and image metadata.
 export declare type Metadata = SystemConfig & {
--- a/frontend/src/app/socketio/emitters.ts
+++ b/frontend/src/app/socketio/emitters.ts
@@ -30,14 +30,14 @@ const makeSocketIOEmitters = (
        options.shouldUseInitImage = false;
      }

-      const { generationParameters, esrganParameters, gfpganParameters } =
+      const { generationParameters, esrganParameters, facetoolParameters } =
        frontendToBackendParameters(options, getState().system);

      socketio.emit(
        'generateImage',
        generationParameters,
        esrganParameters,
-        gfpganParameters
+        facetoolParameters
      );

      dispatch(
@@ -46,7 +46,7 @@ const makeSocketIOEmitters = (
          message: `Image generation requested: ${JSON.stringify({
            ...generationParameters,
            ...esrganParameters,
-            ...gfpganParameters,
+            ...facetoolParameters,
          })}`,
        })
      );
@@ -71,24 +71,32 @@ const makeSocketIOEmitters = (
        })
      );
    },
-    emitRunGFPGAN: (imageToProcess: InvokeAI.Image) => {
+    emitRunFacetool: (imageToProcess: InvokeAI.Image) => {
      dispatch(setIsProcessing(true));
-      const { gfpganStrength } = getState().options;
+      const { facetoolType, facetoolStrength, codeformerFidelity } =
+        getState().options;

-      const gfpganParameters = {
-        facetool_strength: gfpganStrength,
+      const facetoolParameters: Record<string, any> = {
+        facetool_strength: facetoolStrength,
      };
+
+      if (facetoolType === 'codeformer') {
+        facetoolParameters.codeformer_fidelity = codeformerFidelity;
+      }
+
      socketio.emit('runPostprocessing', imageToProcess, {
-        type: 'gfpgan',
-        ...gfpganParameters,
+        type: facetoolType,
+        ...facetoolParameters,
      });
      dispatch(
        addLogEntry({
          timestamp: dateFormat(new Date(), 'isoDateTime'),
-          message: `GFPGAN fix faces requested: ${JSON.stringify({
-            file: imageToProcess.url,
-            ...gfpganParameters,
-          })}`,
+          message: `Face restoration (${facetoolType}) requested: ${JSON.stringify(
+            {
+              file: imageToProcess.url,
+              ...facetoolParameters,
+            }
+          )}`,
        })
      );
    },
--- a/frontend/src/app/socketio/listeners.ts
+++ b/frontend/src/app/socketio/listeners.ts
@@ -151,32 +151,6 @@ const makeSocketIOListeners = (
        console.error(e);
      }
    },
-    /**
-     * Callback to run when we receive a 'gfpganResult' event.
-     */
-    onGFPGANResult: (data: InvokeAI.ImageResultResponse) => {
-      try {
-        const { url, metadata, mtime } = data;
-
-        dispatch(
-          addImage({
-            uuid: uuidv4(),
-            url,
-            mtime,
-            metadata,
-          })
-        );
-
-        dispatch(
-          addLogEntry({
-            timestamp: dateFormat(new Date(), 'isoDateTime'),
-            message: `Fixed faces: ${url}`,
-          })
-        );
-      } catch (e) {
-        console.error(e);
-      }
-    },
    /**
     * Callback to run when we receive a 'progressUpdate' event.
     * TODO: Add additional progress phases
--- a/frontend/src/app/socketio/middleware.ts
+++ b/frontend/src/app/socketio/middleware.ts
@@ -50,7 +50,7 @@ export const socketioMiddleware = () => {
    const {
      emitGenerateImage,
      emitRunESRGAN,
-      emitRunGFPGAN,
+      emitRunFacetool: emitRunGFPGAN,
      emitDeleteImage,
      emitRequestImages,
      emitRequestNewImages,
--- a/frontend/src/common/util/parameterTranslation.ts
+++ b/frontend/src/common/util/parameterTranslation.ts
@@ -41,8 +41,10 @@ export const frontendToBackendParameters = (
    shouldRunESRGAN,
    upscalingLevel,
    upscalingStrength,
-    shouldRunGFPGAN,
-    gfpganStrength,
+    shouldRunFacetool,
+    facetoolStrength,
+    codeformerFidelity,
+    facetoolType,
    shouldRandomizeSeed,
  } = optionsState;

@@ -88,7 +90,7 @@ export const frontendToBackendParameters = (
  }

  let esrganParameters: false | { [k: string]: any } = false;
-  let gfpganParameters: false | { [k: string]: any } = false;
+  let facetoolParameters: false | { [k: string]: any } = false;

  if (shouldRunESRGAN) {
    esrganParameters = {
@@ -97,99 +99,19 @@ export const frontendToBackendParameters = (
    };
  }

-  if (shouldRunGFPGAN) {
-    gfpganParameters = {
-      strength: gfpganStrength,
+  if (shouldRunFacetool) {
+    facetoolParameters = {
+      type: facetoolType,
+      strength: facetoolStrength,
    };
+    if (facetoolType === 'codeformer') {
+      facetoolParameters.codeformer_fidelity = codeformerFidelity
+    }
  }

  return {
    generationParameters,
    esrganParameters,
-    gfpganParameters,
+    facetoolParameters,
  };
 };
-
-export const backendToFrontendParameters = (parameters: {
-  [key: string]: any;
-}) => {
-  const {
-    prompt,
-    iterations,
-    steps,
-    cfg_scale,
-    threshold,
-    perlin,
-    height,
-    width,
-    sampler_name,
-    seed,
-    seamless,
-    hires_fix,
-    progress_images,
-    variation_amount,
-    with_variations,
-    facetool_strength,
-    upscale,
-    init_img,
-    init_mask,
-    strength,
-  } = parameters;
-
-  const options: { [key: string]: any } = {
-    shouldDisplayInProgress: progress_images,
-    // init
-    shouldGenerateVariations: false,
-    shouldRunESRGAN: false,
-    shouldRunGFPGAN: false,
-    initialImagePath: '',
-    maskPath: '',
-  };
-
-  if (variation_amount > 0) {
-    options.shouldGenerateVariations = true;
-    options.variationAmount = variation_amount;
-    if (with_variations) {
-      options.seedWeights = seedWeightsToString(with_variations);
-    }
-  }
-
-  if (facetool_strength > 0) {
-    options.shouldRunGFPGAN = true;
-    options.gfpganStrength = facetool_strength;
-  }
-
-  if (upscale) {
-    options.shouldRunESRGAN = true;
-    options.upscalingLevel = upscale[0];
-    options.upscalingStrength = upscale[1];
-  }
-
-  if (init_img) {
-    options.shouldUseInitImage = true;
-    options.initialImagePath = init_img;
-    options.strength = strength;
-    if (init_mask) {
-      options.maskPath = init_mask;
-    }
-  }
-
-  // if we had a prompt, add all the metadata, but if we don't have a prompt,
-  // we must have only done ESRGAN or GFPGAN so do not add that metadata
-  if (prompt) {
-    options.prompt = prompt;
-    options.iterations = iterations;
-    options.steps = steps;
-    options.cfgScale = cfg_scale;
-    options.threshold = threshold;
-    options.perlin = perlin;
-    options.height = height;
-    options.width = width;
-    options.sampler = sampler_name;
-    options.seed = seed;
-    options.seamless = seamless;
-    options.hiresFix = hires_fix;
-  }
-
-  return options;
-};
--- a/frontend/src/features/gallery/CurrentImageButtons.tsx
+++ b/frontend/src/features/gallery/CurrentImageButtons.tsx
@@ -66,8 +66,8 @@ const CurrentImageButtons = ({ image }: CurrentImageButtonsProps) => {
    (state: RootState) => state.options.upscalingLevel
  );

-  const gfpganStrength = useAppSelector(
-    (state: RootState) => state.options.gfpganStrength
+  const facetoolStrength = useAppSelector(
+    (state: RootState) => state.options.facetoolStrength
  );

  const { isProcessing, isConnected, isGFPGANAvailable, isESRGANAvailable } =
@@ -195,7 +195,7 @@ const CurrentImageButtons = ({ image }: CurrentImageButtonsProps) => {
        Boolean(!intermediateImage) &&
        isConnected &&
        !isProcessing &&
-        gfpganStrength
+        facetoolStrength
      ) {
        handleClickFixFaces();
      } else {
@@ -213,7 +213,7 @@ const CurrentImageButtons = ({ image }: CurrentImageButtonsProps) => {
      intermediateImage,
      isConnected,
      isProcessing,
-      gfpganStrength,
+      facetoolStrength,
    ]
  );

@@ -270,7 +270,7 @@ const CurrentImageButtons = ({ image }: CurrentImageButtonsProps) => {
              !isGFPGANAvailable ||
              Boolean(intermediateImage) ||
              !(isConnected && !isProcessing) ||
-              !gfpganStrength
+              !facetoolStrength
            }
            onClick={handleClickFixFaces}
          />
--- a/frontend/src/features/gallery/ImageMetaDataViewer/ImageMetadataViewer.tsx
+++ b/frontend/src/features/gallery/ImageMetaDataViewer/ImageMetadataViewer.tsx
@@ -14,7 +14,9 @@ import { useAppDispatch } from '../../../app/store';
 import * as InvokeAI from '../../../app/invokeai';
 import {
  setCfgScale,
-  setGfpganStrength,
+  setFacetoolStrength,
+  setCodeformerFidelity,
+  setFacetoolType,
  setHeight,
  setHiresFix,
  setImg2imgStrength,
@@ -151,7 +153,7 @@ const ImageMetadataViewer = memo(
                <MetadataItem
                  label="Fix faces strength"
                  value={strength}
-                  onClick={() => dispatch(setGfpganStrength(strength))}
+                  onClick={() => dispatch(setFacetoolStrength(strength))}
                />
              )}
              {type === 'esrgan' && scale !== undefined && (
@@ -321,12 +323,46 @@ const ImageMetadataViewer = memo(
                            <MetadataItem
                              label="Strength"
                              value={strength}
-                              onClick={() =>
-                                dispatch(setGfpganStrength(strength))
-                              }
+                              onClick={() => {
+                                dispatch(setFacetoolStrength(strength));
+                                dispatch(setFacetoolType('gfpgan'));
+                              }}
                            />
                          </Flex>
                        );
+                      } else if (postprocess.type === 'codeformer') {
+                        const { strength, fidelity } = postprocess;
+                        return (
+                          <Flex
+                            key={i}
+                            pl={'2rem'}
+                            gap={1}
+                            direction={'column'}
+                          >
+                            <Text size={'md'}>{`${
+                              i + 1
+                            }: Face restoration (Codeformer)`}</Text>
+
+                            <MetadataItem
+                              label="Strength"
+                              value={strength}
+                              onClick={() => {
+                                dispatch(setFacetoolStrength(strength));
+                                dispatch(setFacetoolType('codeformer'));
+                              }}
+                            />
+                            {fidelity && (
+                              <MetadataItem
+                                label="Fidelity"
+                                value={fidelity}
+                                onClick={() => {
+                                  dispatch(setCodeformerFidelity(fidelity));
+                                  dispatch(setFacetoolType('codeformer'));
+                                }}
+                              />
+                            )}
+                          </Flex>
+                        );
                      }
                    }
                  )}
--- a/frontend/src/features/options/AdvancedOptions/FaceRestore/FaceRestore.tsx
+++ b/frontend/src/features/options/AdvancedOptions/FaceRestore/FaceRestore.tsx
@@ -6,21 +6,21 @@ import {
  useAppSelector,
 } from '../../../../app/store';
 import IAISwitch from '../../../../common/components/IAISwitch';
-import { setShouldRunGFPGAN } from '../../optionsSlice';
+import { setShouldRunFacetool } from '../../optionsSlice';

 export default function FaceRestore() {
  const isGFPGANAvailable = useAppSelector(
    (state: RootState) => state.system.isGFPGANAvailable
  );

-  const shouldRunGFPGAN = useAppSelector(
-    (state: RootState) => state.options.shouldRunGFPGAN
+  const shouldRunFacetool = useAppSelector(
+    (state: RootState) => state.options.shouldRunFacetool
  );

  const dispatch = useAppDispatch();

-  const handleChangeShouldRunGFPGAN = (e: ChangeEvent<HTMLInputElement>) =>
-    dispatch(setShouldRunGFPGAN(e.target.checked));
+  const handleChangeShouldRunFacetool = (e: ChangeEvent<HTMLInputElement>) =>
+    dispatch(setShouldRunFacetool(e.target.checked));

  return (
    <Flex
@@ -32,8 +32,8 @@ export default function FaceRestore() {
      <p>Restore Face</p>
      <IAISwitch
        isDisabled={!isGFPGANAvailable}
-        isChecked={shouldRunGFPGAN}
-        onChange={handleChangeShouldRunGFPGAN}
+        isChecked={shouldRunFacetool}
+        onChange={handleChangeShouldRunFacetool}
      />
    </Flex>
  );
--- a/frontend/src/features/options/AdvancedOptions/FaceRestore/FaceRestoreOptions.tsx
+++ b/frontend/src/features/options/AdvancedOptions/FaceRestore/FaceRestoreOptions.tsx
@@ -3,18 +3,29 @@ import { Flex } from '@chakra-ui/react';
 import { RootState } from '../../../../app/store';
 import { useAppDispatch, useAppSelector } from '../../../../app/store';

-import { OptionsState, setGfpganStrength } from '../../optionsSlice';
+import {
+  FacetoolType,
+  OptionsState,
+  setCodeformerFidelity,
+  setFacetoolStrength,
+  setFacetoolType,
+} from '../../optionsSlice';

 import { createSelector } from '@reduxjs/toolkit';
 import { isEqual } from 'lodash';
 import { SystemState } from '../../../system/systemSlice';
 import IAINumberInput from '../../../../common/components/IAINumberInput';
+import IAISelect from '../../../../common/components/IAISelect';
+import { FACETOOL_TYPES } from '../../../../app/constants';
+import { ChangeEvent } from 'react';

 const optionsSelector = createSelector(
  (state: RootState) => state.options,
  (options: OptionsState) => {
    return {
-      gfpganStrength: options.gfpganStrength,
+      facetoolStrength: options.facetoolStrength,
+      facetoolType: options.facetoolType,
+      codeformerFidelity: options.codeformerFidelity,
    };
  },
  {
@@ -43,13 +54,26 @@ const systemSelector = createSelector(
 */
 const FaceRestoreOptions = () => {
  const dispatch = useAppDispatch();
-  const { gfpganStrength } = useAppSelector(optionsSelector);
+  const { facetoolStrength, facetoolType, codeformerFidelity } =
+    useAppSelector(optionsSelector);
  const { isGFPGANAvailable } = useAppSelector(systemSelector);

-  const handleChangeStrength = (v: number) => dispatch(setGfpganStrength(v));
+  const handleChangeStrength = (v: number) => dispatch(setFacetoolStrength(v));
+
+  const handleChangeCodeformerFidelity = (v: number) =>
+    dispatch(setCodeformerFidelity(v));
+
+  const handleChangeFacetoolType = (e: ChangeEvent<HTMLSelectElement>) =>
+    dispatch(setFacetoolType(e.target.value as FacetoolType));

  return (
    <Flex direction={'column'} gap={2}>
+      <IAISelect
+        label="Type"
+        validValues={FACETOOL_TYPES.concat()}
+        value={facetoolType}
+        onChange={handleChangeFacetoolType}
+      />
      <IAINumberInput
        isDisabled={!isGFPGANAvailable}
        label="Strength"
@@ -57,10 +81,23 @@ const FaceRestoreOptions = () => {
        min={0}
        max={1}
        onChange={handleChangeStrength}
-        value={gfpganStrength}
+        value={facetoolStrength}
        width="90px"
        isInteger={false}
      />
+      {facetoolType === 'codeformer' && (
+        <IAINumberInput
+          isDisabled={!isGFPGANAvailable}
+          label="Fidelity"
+          step={0.05}
+          min={0}
+          max={1}
+          onChange={handleChangeCodeformerFidelity}
+          value={codeformerFidelity}
+          width="90px"
+          isInteger={false}
+        />
+      )}
    </Flex>
  );
 };
--- a/frontend/src/features/options/optionsSlice.ts
+++ b/frontend/src/features/options/optionsSlice.ts
@@ -3,9 +3,12 @@ import type { PayloadAction } from '@reduxjs/toolkit';
 import * as InvokeAI from '../../app/invokeai';
 import promptToString from '../../common/util/promptToString';
 import { seedWeightsToString } from '../../common/util/seedWeightPairs';
+import { FACETOOL_TYPES } from '../../app/constants';

 export type UpscalingLevel = 2 | 4;

+export type FacetoolType = typeof FACETOOL_TYPES[number];
+
 export interface OptionsState {
  prompt: string;
  iterations: number;
@@ -18,7 +21,9 @@ export interface OptionsState {
  perlin: number;
  seed: number;
  img2imgStrength: number;
-  gfpganStrength: number;
+  facetoolType: FacetoolType;
+  facetoolStrength: number;
+  codeformerFidelity: number;
  upscalingLevel: UpscalingLevel;
  upscalingStrength: number;
  shouldUseInitImage: boolean;
@@ -31,7 +36,7 @@ export interface OptionsState {
  variationAmount: number;
  seedWeights: string;
  shouldRunESRGAN: boolean;
-  shouldRunGFPGAN: boolean;
+  shouldRunFacetool: boolean;
  shouldRandomizeSeed: boolean;
  showAdvancedOptions: boolean;
  activeTab: number;
@@ -63,8 +68,10 @@ const initialOptionsState: OptionsState = {
  shouldRunESRGAN: false,
  upscalingLevel: 4,
  upscalingStrength: 0.75,
-  shouldRunGFPGAN: false,
-  gfpganStrength: 0.8,
+  shouldRunFacetool: false,
+  facetoolStrength: 0.8,
+  facetoolType: 'gfpgan',
+  codeformerFidelity: 0.75,
  shouldRandomizeSeed: true,
  showAdvancedOptions: true,
  activeTab: 0,
@@ -117,8 +124,11 @@ export const optionsSlice = createSlice({
    setImg2imgStrength: (state, action: PayloadAction<number>) => {
      state.img2imgStrength = action.payload;
    },
-    setGfpganStrength: (state, action: PayloadAction<number>) => {
-      state.gfpganStrength = action.payload;
+    setFacetoolStrength: (state, action: PayloadAction<number>) => {
+      state.facetoolStrength = action.payload;
+    },
+    setCodeformerFidelity: (state, action: PayloadAction<number>) => {
+      state.codeformerFidelity = action.payload;
    },
    setUpscalingLevel: (state, action: PayloadAction<UpscalingLevel>) => {
      state.upscalingLevel = action.payload;
@@ -229,8 +239,8 @@ export const optionsSlice = createSlice({
      //     (postprocess: InvokeAI.PostProcessedImageMetadata) => {
      //       if (postprocess.type === 'gfpgan') {
      //         const { strength } = postprocess;
-      //         if (strength) state.gfpganStrength = strength;
-      //         state.shouldRunGFPGAN = true;
+      //         if (strength) state.facetoolStrength = strength;
+      //         state.shouldRunFacetool = true;
      //         postprocessingNotDone = postprocessingNotDone.filter(
      //           (p) => p !== 'gfpgan'
      //         );
@@ -250,7 +260,7 @@ export const optionsSlice = createSlice({

      // postprocessingNotDone.forEach((p) => {
      //   if (p === 'esrgan') state.shouldRunESRGAN = false;
-      //   if (p === 'gfpgan') state.shouldRunGFPGAN = false;
+      //   if (p === 'gfpgan') state.shouldRunFacetool = false;
      // });

      if (prompt) state.prompt = promptToString(prompt);
@@ -260,7 +270,7 @@ export const optionsSlice = createSlice({
      if (threshold) state.threshold = threshold;
      if (typeof threshold === 'undefined') state.threshold = 0;
      if (perlin) state.perlin = perlin;
-      if (typeof perlin === 'undefined') state.perlin = 0;      
+      if (typeof perlin === 'undefined') state.perlin = 0;
      if (typeof seamless === 'boolean') state.seamless = seamless;
      if (typeof hires_fix === 'boolean') state.hiresFix = hires_fix;
      if (width) state.width = width;
@@ -272,8 +282,11 @@ export const optionsSlice = createSlice({
        ...initialOptionsState,
      };
    },
-    setShouldRunGFPGAN: (state, action: PayloadAction<boolean>) => {
-      state.shouldRunGFPGAN = action.payload;
+    setShouldRunFacetool: (state, action: PayloadAction<boolean>) => {
+      state.shouldRunFacetool = action.payload;
+    },
+    setFacetoolType: (state, action: PayloadAction<FacetoolType>) => {
+      state.facetoolType = action.payload;
    },
    setShouldRunESRGAN: (state, action: PayloadAction<boolean>) => {
      state.shouldRunESRGAN = action.payload;
@@ -310,7 +323,9 @@ export const {
  setSeamless,
  setHiresFix,
  setImg2imgStrength,
-  setGfpganStrength,
+  setFacetoolStrength,
+  setFacetoolType,
+  setCodeformerFidelity,
  setUpscalingLevel,
  setUpscalingStrength,
  setShouldUseInitImage,
@@ -324,7 +339,7 @@ export const {
  setSeedWeights,
  setVariationAmount,
  setAllParameters,
-  setShouldRunGFPGAN,
+  setShouldRunFacetool,
  setShouldRunESRGAN,
  setShouldRandomizeSeed,
  setShowAdvancedOptions,