fix: send missing model params, add face strength control

ssube · Jan 17, 2023 · 0e27cc8 · 0e27cc8
1 parent 83f2e87
commit 0e27cc8
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 13 deletions.
diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py
@@ -183,6 +183,7 @@ def upscale_from_request() -> UpscaleParams:
     upscaling = get_from_list(request.args, 'upscaling', upscaling_models)
     correction = get_from_list(request.args, 'correction', correction_models)
     faces = request.args.get('faces', 'false') == 'true'
+    face_strength = get_and_clamp_float(request.args, 'faceStrength', 0.5, 1.0, 0.0)
 
     return UpscaleParams(
         upscaling,
@@ -192,6 +193,7 @@ def upscale_from_request() -> UpscaleParams:
         faces=faces,
         platform='onnx',
         denoise=denoise,
+        face_strength=face_strength,
     )
 
 

diff --git a/api/onnx_web/upscale.py b/api/onnx_web/upscale.py
@@ -91,6 +91,7 @@ def __init__(
         outscale: int = 1,
         denoise: float = 0.5,
         faces=True,
+        face_strength: float = 0.5,
         platform: str = 'onnx',
         half=False
     ) -> None:
@@ -100,21 +101,10 @@ def __init__(
         self.outscale = outscale
         self.denoise = denoise
         self.faces = faces
+        self.face_strength = face_strength
         self.platform = platform
         self.half = half
 
-    def rescale(self, scale: int, outscale: int = 1):
-        return UpscaleParams(
-            self.upscale_model,
-            correction_model=self.correction_model,
-            scale=scale,
-            outscale=outscale,
-            denoise=self.denoise,
-            faces=self.faces,
-            platform=self.platform,
-            half=self.half,
-        )
-
     def resize(self, size: Size) -> Size:
         return Size(size.width * self.scale * self.outscale, size.height * self.scale * self.outscale)
 
@@ -190,6 +180,6 @@ def upscale_gfpgan(ctx: ServerContext, params: UpscaleParams, image, upsampler=N
         bg_upsampler=upsampler)
 
     _, _, output = face_enhancer.enhance(
-        image, has_aligned=False, only_center_face=False, paste_back=True)
+        image, has_aligned=False, only_center_face=False, paste_back=True, weight=params.face_strength)
 
     return output
diff --git a/gui/src/client.ts b/gui/src/client.ts
@@ -73,6 +73,7 @@ export interface UpscaleParams {
   faces: boolean;
   scale: number;
   outscale: number;
+  faceStrength: number;
 }
 
 export interface ApiResponse {
@@ -167,6 +168,8 @@ export function makeImageURL(root: string, type: string, params: BaseImgParams):
 export function appendModelToURL(url: URL, params: ModelParams) {
   url.searchParams.append('model', params.model);
   url.searchParams.append('platform', params.platform);
+  url.searchParams.append('upscaling', params.upscaling);
+  url.searchParams.append('correction', params.correction);
 }
 
 export function appendUpscaleToURL(url: URL, upscale: UpscaleParams) {
@@ -175,6 +178,7 @@ export function appendUpscaleToURL(url: URL, upscale: UpscaleParams) {
     url.searchParams.append('faces', String(upscale.faces));
     url.searchParams.append('scale', upscale.scale.toFixed(FIXED_INTEGER));
     url.searchParams.append('outscale', upscale.outscale.toFixed(FIXED_INTEGER));
+    url.searchParams.append('faceStrength', upscale.faceStrength.toFixed(FIXED_FLOAT));
   }
 }
 

diff --git a/gui/src/components/UpscaleControl.tsx b/gui/src/components/UpscaleControl.tsx
@@ -89,5 +89,19 @@ export function UpscaleControl(props: UpscaleControlProps) {
       <FaceRetouchingNatural />
       Face Correction
     </ToggleButton>
+    <NumericField
+      label='Strength'
+      decimal
+      disabled={params.enabled === false}
+      min={config.faceStrength.min}
+      max={config.faceStrength.max}
+      step={config.faceStrength.step}
+      value={params.faceStrength}
+      onChange={(faceStrength) => {
+        setUpscale({
+          faceStrength,
+        });
+      }}
+    />
   </Stack>;
 }
diff --git a/gui/src/state.ts b/gui/src/state.ts
@@ -250,6 +250,7 @@ export function createStateSlices(base: ConfigParams) {
       faces: false,
       scale: 1,
       outscale: 1,
+      faceStrength: 0.5,
     },
     setUpscale(upscale) {
       set((prev) => ({