Merge pull request #74 from freek99/feat/upscale-customx

Add upscale custom support

Author: konieshadow

fooocusapi/api_utils.py

@@ -47,6 +47,8 @@ def req_to_params(req: Text2ImgRequest) -> ImageGenerationParams:
         req, ImgUpscaleOrVaryRequest) or isinstance(req, ImgUpscaleOrVaryRequestJson)) else read_input_image(req.input_image)
     uov_method = flags.disabled if not (isinstance(
         req, ImgUpscaleOrVaryRequest) or isinstance(req, ImgUpscaleOrVaryRequestJson)) else req.uov_method.value
+    upscale_value = None if not (isinstance(
+                req, ImgUpscaleOrVaryRequest) or isinstance(req, ImgUpscaleOrVaryRequestJson)) else req.upscale_value
     outpaint_selections = [] if not (isinstance(
         req, ImgInpaintOrOutpaintRequest) or isinstance(req, ImgInpaintOrOutpaintRequestJson)) else [
         s.value for s in req.outpaint_selections]
@@ -133,6 +135,7 @@ def req_to_params(req: Text2ImgRequest) -> ImageGenerationParams:
                                  loras=loras,
                                  uov_input_image=uov_input_image,
                                  uov_method=uov_method,
+                                 upscale_value=upscale_value,
                                  outpaint_selections=outpaint_selections,
                                  outpaint_distance_left=outpaint_distance_left,
                                  outpaint_distance_right=outpaint_distance_right,

fooocusapi/models.py

@@ -41,7 +41,7 @@ class UpscaleOrVaryMethod(str, Enum):
     upscale_15 = 'Upscale (1.5x)'
     upscale_2 = 'Upscale (2x)'
     upscale_fast = 'Upscale (Fast 2x)'
-
+    upscale_custom = 'Upscale (Custom)'
 
 class OutpaintExpansion(str, Enum):
     left = 'Left'
@@ -122,10 +122,12 @@ class Text2ImgRequest(BaseModel):
 class ImgUpscaleOrVaryRequest(Text2ImgRequest):
     input_image: UploadFile
     uov_method: UpscaleOrVaryMethod
+    upscale_value: float
 
     @classmethod
     def as_form(cls, input_image: UploadFile = Form(description="Init image for upsacale or outpaint"),
                 uov_method: UpscaleOrVaryMethod = Form(),
+                upscale_value: float = Form(default=1.0, ge=1.0, le=5.0),
                 prompt: str = Form(''),
                 negative_prompt: str = Form(default_prompt_negative),
                 style_selections: List[str] = Form(defualt_styles, description="Fooocus style selections, seperated by comma"),
@@ -168,7 +170,7 @@ def as_form(cls, input_image: UploadFile = Form(description="Init image for upsa
                 errs = ve.errors()
                 raise RequestValidationError(errors=[errs])
 
-        return cls(input_image=input_image, uov_method=uov_method, prompt=prompt, negative_prompt=negative_prompt, style_selections=style_selection_arr,
+        return cls(input_image=input_image, uov_method=uov_method,upscale_value=upscale_value, prompt=prompt, negative_prompt=negative_prompt, style_selections=style_selection_arr,
                    performance_selection=performance_selection, aspect_ratios_selection=aspect_ratios_selection,
                    image_number=image_number, image_seed=image_seed, sharpness=sharpness, guidance_scale=guidance_scale,
                    base_model_name=base_model_name, refiner_model_name=refiner_model_name, refiner_switch=refiner_switch,

fooocusapi/models_v2.py

@@ -3,6 +3,7 @@
 
 class ImgUpscaleOrVaryRequestJson(Text2ImgRequest):
     uov_method: UpscaleOrVaryMethod = "Upscale (2x)"
+    upscale_value: float | None = Field(None, ge=1.0, le=5.0, description="Upscale custom value, None for default value")
     input_image: str = Field(description="Init image for upsacale or outpaint as base64")
 
 

fooocusapi/parameters.py

@@ -93,6 +93,7 @@ def __init__(self, prompt: str,
                  loras: List[Tuple[str, float]],
                  uov_input_image: np.ndarray | None,
                  uov_method: str,
+                 upscale_value: float,
                  outpaint_selections: List[str],
                  outpaint_distance_left: int,
                  outpaint_distance_right: int,
@@ -117,6 +118,7 @@ def __init__(self, prompt: str,
         self.loras = loras
         self.uov_input_image = uov_input_image
         self.uov_method = uov_method
+        self.upscale_value = upscale_value
         self.outpaint_selections = outpaint_selections
         self.outpaint_distance_left = outpaint_distance_left
         self.outpaint_distance_right = outpaint_distance_right

fooocusapi/worker.py

@@ -3,6 +3,7 @@
 import time
 import numpy as np
 import torch
+import re
 from typing import List
 from fooocusapi.file_utils import save_output_file
 from fooocusapi.parameters import GenerationFinishReason, ImageGenerationParams, ImageGenerationResult
@@ -121,6 +122,7 @@ def yield_result(_, imgs, tasks):
         input_image_checkbox = params.uov_input_image is not None or params.inpaint_input_image is not None or len(params.image_prompts) > 0
         current_tab = 'uov' if params.uov_method != flags.disabled else 'inpaint' if params.inpaint_input_image is not None else 'ip' if len(params.image_prompts) > 0 else None
         uov_method = params.uov_method
+        upscale_value = params.upscale_value
         uov_input_image = params.uov_input_image
         outpaint_selections = params.outpaint_selections
         outpaint_distance_left = params.outpaint_distance_left
@@ -466,12 +468,15 @@ def yield_result(_, imgs, tasks):
             uov_input_image = perform_upscale(uov_input_image)
             print(f'Image upscaled.')
 
-            if '1.5x' in uov_method:
-                f = 1.5
-            elif '2x' in uov_method:
-                f = 2.0
+            f = 1.0
+            if upscale_value > 1.0:
+                f = upscale_value
             else:
-                f = 1.0
+                pattern = r"([0-9]+(?:\.[0-9]+)?)x"
+                matches = re.findall(pattern, uov_method)
+                if len(matches) > 0:
+                    f_tmp = float(matches[0])
+                    f = 1.0 if f_tmp < 1.0 else 5.0 if f_tmp > 5.0 else f_tmp
 
             shape_ceil = get_shape_ceil(H * f, W * f)