Upload folder using huggingface_hub

inference2.py

@@ -7,18 +7,17 @@ from internals.data.task import ModelType, Task, TaskType
 from internals.pipelines.inpainter import InPainter
 from internals.pipelines.object_remove import ObjectRemoval
 from internals.pipelines.prompt_modifier import PromptModifier
-from internals.pipelines.remove_background import RemoveBackground, RemoveBackgroundV2
+from internals.pipelines.remove_background import (RemoveBackground,
+                                                   RemoveBackgroundV2)
 from internals.pipelines.replace_background import ReplaceBackground
 from internals.pipelines.safety_checker import SafetyChecker
 from internals.pipelines.upscaler import Upscaler
 from internals.util.avatar import Avatar
 from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda
-from internals.util.commons import construct_default_s3_url, upload_image, upload_images
-from internals.util.config import (
-    num_return_sequences,
-    set_configs_from_task,
-    set_root_dir,
-)
+from internals.util.commons import (construct_default_s3_url, upload_image,
+                                    upload_images)
+from internals.util.config import (num_return_sequences, set_configs_from_task,
+                                   set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.slack import Slack
 
@@ -135,12 +134,20 @@ def upscale_image(task: Task):
     if task.get_modelType() == ModelType.ANIME:
         print("Using Anime model")
         out_img = upscaler.upscale_anime(
-            image=task.get_imageUrl(), resize_dimension=task.get_resize_dimension()
+            image=task.get_imageUrl(),
+            width=task.get_width(),
+            height=task.get_height(),
+            face_enhance=task.get_face_enhance(),
+            resize_dimension=task.get_resize_dimension(),
         )
     else:
         print("Using Real model")
         out_img = upscaler.upscale(
-            image=task.get_imageUrl(), resize_dimension=task.get_resize_dimension()
+            image=task.get_imageUrl(),
+            width=task.get_width(),
+            height=task.get_height(),
+            face_enhance=task.get_face_enhance(),
+            resize_dimension=task.get_resize_dimension(),
         )
 
     upload_image(BytesIO(out_img), output_key)

internals/data/task.py

@@ -112,6 +112,9 @@ class Task:
     def get_resize_dimension(self) -> int:
         return self.__data.get("resize_dimension", 1024)
 
+    def get_face_enhance(self) -> bool:
+        return self.__data.get("up_face_enhance", False)
+
     def get_ti_guidance_scale(self) -> float:
         return self.__data.get("ti_guidance_scale", 7.5)
 

internals/pipelines/replace_background.py

@@ -106,7 +106,14 @@ class ReplaceBackground:
         if not has_nsfw:
             for i in range(len(images)):
                 images[i].paste(image, (0, 0), image)
-                out_bytes = self.upscaler.upscale(images[i], resize_dimension)
+                w, h = images[i].size
+                out_bytes = self.upscaler.upscale(
+                    image=images[i],
+                    width=w,
+                    height=h,
+                    face_enhance=False,
+                    resize_dimension=resize_dimension,
+                )
                 images[i] = Image.open(BytesIO(out_bytes)).convert("RGB")
 
         return (images, has_nsfw)

internals/pipelines/upscaler.py

@@ -6,6 +6,7 @@ from typing import Union
 import cv2
 import numpy as np
 from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
 from basicsr.utils.download_util import load_file_from_url
 from gfpgan import GFPGANer
 from PIL import Image
@@ -17,7 +18,7 @@ from internals.util.config import get_root_dir
 
 
 class Upscaler:
-    __model_esrgan_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+    __model_esrgan_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
     __model_esrgan_anime_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
     __model_gfpgan_url = (
         "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
@@ -41,21 +42,39 @@ class Upscaler:
         )
         self.__loaded = True
 
-    def upscale(self, image: Union[str, Image.Image], resize_dimension: int) -> bytes:
-        model = RRDBNet(
+    def upscale(
+        self,
+        image: Union[str, Image.Image],
+        width: int,
+        height: int,
+        face_enhance: bool,
+        resize_dimension: int,
+    ) -> bytes:
+        model = SRVGGNetCompact(
             num_in_ch=3,
             num_out_ch=3,
             num_feat=64,
-            num_block=23,
-            num_grow_ch=32,
-            scale=4,
+            num_conv=32,
+            upscale=4,
+            act_type="prelu",
         )
         return self.__internal_upscale(
-            image, resize_dimension, self.__model_path, model
+            image,
+            resize_dimension,
+            face_enhance,
+            width,
+            height,
+            self.__model_path,
+            model,
         )
 
     def upscale_anime(
-        self, image: Union[str, Image.Image], resize_dimension: int
+        self,
+        image: Union[str, Image.Image],
+        width: int,
+        height: int,
+        face_enhance: bool,
+        resize_dimension: int,
     ) -> bytes:
         model = RRDBNet(
             num_in_ch=3,
@@ -66,7 +85,13 @@ class Upscaler:
             scale=4,
         )
         return self.__internal_upscale(
-            image, resize_dimension, self.__model_path_anime, model
+            image,
+            resize_dimension,
+            face_enhance,
+            width,
+            height,
+            self.__model_path_anime,
+            model,
         )
 
     def __preload_model(self, url: str, download_dir: Path):
@@ -85,25 +110,34 @@ class Upscaler:
         self,
         image,
         resize_dimension: int,
+        face_enhance: bool,
+        width: int,
+        height: int,
         model_path: str,
-        rrbdnet: RRDBNet,
+        model,
     ) -> bytes:
         if type(image) is str:
             image = download_image(image)
-            image = ImageUtil.resize_image_to512(image)
-            image = ImageUtil.to_bytes(image)
+            w, h = image.size
+            if max(w, h) > 1536:
+                image = ImageUtil.resize_image(image, dimension=1536)
 
-        if isinstance(image, Image.Image):
-            image = ImageUtil.to_bytes(image)
-
-        image_array = np.frombuffer(image, dtype=np.uint8)
-        input_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
+        in_path = str(Path.home() / ".cache" / "input_upscale.png")
+        image.save(in_path)
+        input_image = cv2.imread(in_path, cv2.IMREAD_UNCHANGED)
         dimension = min(input_image.shape[0], input_image.shape[1])
         scale = max(math.floor(resize_dimension / dimension), 2)
 
         os.chdir(str(Path.home() / ".cache"))
         upsampler = RealESRGANer(
-            scale=4, model_path=model_path, model=rrbdnet, half="fp16", gpu_id="0"
+            scale=4,
+            model_path=model_path,
+            model=model,
+            half=False,
+            gpu_id="0",
+            tile=0,
+            tile_pad=10,
+            pre_pad=0,
         )
         face_enhancer = GFPGANer(
             model_path=self.__model_path_gfpgan,
@@ -113,9 +147,13 @@ class Upscaler:
             bg_upsampler=upsampler,
         )
 
-        _, _, output = face_enhancer.enhance(
-            input_image, has_aligned=False, only_center_face=True, paste_back=True
-        )
+        if face_enhance:
+            _, _, output = face_enhancer.enhance(
+                input_image, has_aligned=False, only_center_face=False, paste_back=True
+            )
+        else:
+            output, _ = upsampler.enhance(input_image, outscale=scale)
         os.chdir(get_root_dir())
+        cv2.imwrite("out.png", output)
         out_bytes = cv2.imencode(".png", output)[1].tobytes()
         return out_bytes

internals/util/image.py

@@ -9,12 +9,12 @@ def to_bytes(image: Image.Image) -> bytes:
         return output.getvalue()
 
 
-def resize_image_to512(image: Image.Image) -> Image.Image:
+def resize_image(image: Image.Image, dimension: int = 512) -> Image.Image:
     iw, ih = image.size
     if iw > ih:
-        image = image.resize((512, int(512 * ih / iw)))
+        image = image.resize((dimension, int(dimension * ih / iw)))
     else:
-        image = image.resize((int(512 * iw / ih), 512))
+        image = image.resize((int(dimension * iw / ih), dimension))
     return image
 
 

internals/util/lora_style.py

@@ -176,5 +176,4 @@ class LoraStyle:
 
     @staticmethod
     def unload_lora_weights(pipe):
-        pipe.unet.set_attn_processor(AttnProcessor2_0())  # for pytorch 2.0
-        pipe._remove_text_encoder_monkey_patch()
+        pipe.unload_lora_weights()