Apply superresolution with Real-ESRGAN

app.py

@@ -60,6 +60,8 @@ def create_advanced_demo(model: Model) -> gr.Blocks:
                                      step=1,
                                      value=1234,
                                      label='Seed')
+                    superresolve = gr.Checkbox(value=False,
+                                               label='Superresolve')
                     run_button = gr.Button('Run')
             with gr.Column():
                 with gr.Tabs():
@@ -80,6 +82,7 @@ def create_advanced_demo(model: Model) -> gr.Blocks:
                              num_steps,
                              randomize_seed,
                              seed,
+                             superresolve,
                          ],
                          outputs=[
                              result,

model.py

@@ -6,6 +6,7 @@ import random
 import sys
 import tempfile
 
+import gradio as gr
 import imageio
 import numpy as np
 import PIL.Image
@@ -44,6 +45,8 @@ class Model:
                                             self.scheduler_type)
         self.rng = random.Random()
 
+        self.real_esrgan = gr.Interface.load('spaces/hysts/Real-ESRGAN-anime')
+
     @staticmethod
     def _load_pipeline(model_name: str,
                        scheduler_type: str) -> DiffusionPipeline:
@@ -140,17 +143,29 @@ class Model:
         writer.close()
 
         logger.info('--- done ---')
-        return res, out_file.name
+        return PIL.Image.fromarray(res), out_file.name
+
+    def superresolve(self, image: PIL.Image.Image) -> PIL.Image.Image:
+        logger.info('--- superresolve ---')
+
+        with tempfile.NamedTemporaryFile(suffix='.png') as f:
+            image.save(f.name)
+            out_file = self.real_esrgan(f.name)
+
+        logger.info('--- done ---')
+        return PIL.Image.open(out_file)
 
     def run(self, model_name: str, scheduler_type: str, num_steps: int,
-            randomize_seed: bool,
-            seed: int) -> tuple[PIL.Image.Image, int, str]:
+            randomize_seed: bool, seed: int,
+            superresolve: bool) -> tuple[PIL.Image.Image, int, str]:
         self.set_pipeline(model_name, scheduler_type)
         if scheduler_type == 'PNDM':
             num_steps = max(4, min(num_steps, 100))
         if randomize_seed:
             seed = self.rng.randint(0, 100000)
         res, filename = self.generate_with_video(seed, num_steps)
+        if superresolve:
+            res = self.superresolve(res)
         return res, seed, filename
 
     @staticmethod
@@ -169,4 +184,5 @@ class Model:
         self.set_pipeline(self.MODEL_NAMES[0], 'DDIM')
         seed = self.rng.randint(0, 1000000)
         images = self.generate(seed, num_steps=10, num_images=4)
+        images = [self.superresolve(image) for image in images]
         return self.to_grid(images, 2)