Add Real-ESRGAN

discord_tron_client/classes/image_manipulation/diffusion.py

@@ -12,13 +12,14 @@
 from typing import Dict
 from discord_tron_client.classes.hardware import HardwareInfo
 from discord_tron_client.classes.app_config import AppConfig
+from discord_tron_client.classes.image_manipulation.face_upscale import get_upscaler, use_upscaler
 from PIL import Image
 import torch, gc, logging, diffusers
-# torch.backends.cudnn.deterministic = True
-# torch.backends.cuda.matmul.allow_tf32 = True
-# torch.backends.cudnn.allow_tf32 = True
-# torch.backends.cudnn.benchmark = True
-# torch.backends.cuda.enable_flash_sdp(True)
+torch.backends.cudnn.deterministic = True
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.enable_flash_sdp(True)
 if torch.backends.cuda.mem_efficient_sdp_enabled():
     logging.info("CUDA SDP (scaled dot product attention) is enabled.")
 if torch.backends.cuda.math_sdp_enabled():
@@ -114,6 +115,15 @@ def create_pipeline(self, model_id: str, pipe_type: str) -> Pipeline:
             pipeline.safety_checker = lambda images, clip_input: (images, False)
         return pipeline
 
+    def upscale_image(self, image: Image):
+        self._initialize_upscaler_pipe()
+        return use_upscaler(self.pipelines["upscaler"], image)
+
+    def _initialize_upscaler_pipe(self):
+        if "upscaler" not in self.pipelines:
+            self.pipelines["upscaler"] = get_upscaler()
+        return self.pipelines["upscaler"]
+
     def get_pipe(
         self,
         user_config: dict,

discord_tron_client/classes/image_manipulation/face_upscale.py

@@ -0,0 +1,18 @@
+import torch, os
+from PIL import Image
+import numpy as np
+from RealESRGAN import RealESRGAN
+from discord_tron_client.classes.app_config import AppConfig
+config = AppConfig()
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def get_upscaler(scale: int = 4):
+    model_path = config.get_huggingface_model_path()
+    model = RealESRGAN(device, scale=4)
+    model.load_weights(os.path.join(model_path, 'RealESRGAN_x4.pth'), download=True)
+    return model
+
+def use_upscaler(model: RealESRGAN, image: Image):
+    sr_image = model.predict(image)
+    return sr_image

discord_tron_client/classes/image_manipulation/pipeline.py

@@ -280,6 +280,12 @@ def _run_pipeline(
             except Exception as e:
                 logging.warn(f'Could not cleanly clear the GC: {e}')
 
+        # Now we upscale using Real-ESRGAN.
+        should_upscale = user_config.get('hires_fix', False)
+        if should_upscale:
+            logging.info('Upscaling image using Real-ESRGAN!')
+            new_image = self.pipeline_manager.upscale_image(new_image)
+
         return new_image
 
     async def generate_image(

pyproject.toml

@@ -44,6 +44,7 @@ controlnet-aux = "^0.0.3"
 k-diffusion = "^0.0.14"
 safetensors = "^0.3.1"
 split-image = "^2.0.1"
+realesrgan = {git = "https://github.com/sberbank-ai/Real-ESRGAN.git"}
 
 [[tool.poetry.source]]
 name = "default"