Add model selector

app.py

@@ -13,23 +13,29 @@ from huggingface_hub import hf_hub_download
 from model import build_thera
 from super_resolve import process
 
-REPO_ID = "prs-eth/thera-edsr-plus"
+REPO_ID_EDSR = "prs-eth/thera-edsr-pro"
+REPO_ID_RDN = "prs-eth/thera-rdn-pro"
 
 print(f"JAX devices: {jax.devices()}")
 print(f"JAX device type: {jax.devices()[0].device_kind}")
 
-# load model
-model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pkl")
+model_path = hf_hub_download(repo_id=REPO_ID_EDSR, filename="model.pkl")
 with open(model_path, 'rb') as fh:
     check = pickle.load(fh)
-    params, backbone, size = check['model'], check['backbone'], check['size']
+    params_edsr, backbone, size = check['model'], check['backbone'], check['size']
+    model_edsr = build_thera(3, backbone, size)
 
-model = build_thera(3, backbone, size)
+model_path = hf_hub_download(repo_id=REPO_ID_RDN, filename="model.pkl")
+with open(model_path, 'rb') as fh:
+    check = pickle.load(fh)
+    params_rdn, backbone, size = check['model'], check['backbone'], check['size']
+    model_rdn = build_thera(3, backbone, size)
 
 
 class TheraApp(DualVisionApp):
     DEFAULT_SCALE = 3.92
     DEFAULT_DO_ENSEMBLE = False
+    DEFAULT_MODEL = 'edsr'
 
     def make_header(self):
         gr.Markdown(
@@ -47,7 +53,7 @@ class TheraApp(DualVisionApp):
             </a>
             </p>                    
             <p align="center" style="margin-top: 0px;">
-                Upload a photo or select an example below to do arbitrary-scale super-resolution in real time!
+                <strong>Upload a photo or select an example below to do arbitrary-scale super-resolution in real time!</strong>
             </p>
         """
         )
@@ -61,6 +67,14 @@ class TheraApp(DualVisionApp):
                 step=0.01,
                 value=self.DEFAULT_SCALE,
             )
+            model = gr.Radio(
+                [
+                    ("EDSR", 'edsr'),
+                    ("RDN", 'rdn'),
+                ],
+                label="Backbone",
+                value=self.DEFAULT_MODEL,
+            )
             do_ensemble = gr.Radio(
                 [
                     ("No", False),
@@ -71,12 +85,14 @@ class TheraApp(DualVisionApp):
             )
         return {
             "scale": scale,
+            "model": model,
             "do_ensemble": do_ensemble,
         }
 
     def process(self, image_in: Image.Image, **kwargs):
         scale = kwargs.get("scale", self.DEFAULT_SCALE)
         do_ensemble = kwargs.get("do_ensemble", self.DEFAULT_DO_ENSEMBLE)
+        model = kwargs.get("model", self.DEFAULT_MODEL)
 
         source = np.asarray(image_in) / 255.
 
@@ -86,7 +102,14 @@ class TheraApp(DualVisionApp):
             round(source.shape[1] * scale),
         )
 
-        out = process(source, model, params, target_shape, do_ensemble=do_ensemble)
+        if model == 'edsr':
+            m, p = model_edsr, params_edsr
+        elif model == 'rdn':
+            m, p = model_rdn, params_rdn
+        else:
+            raise NotImplementedError('model:', model)
+
+        out = process(source, m, p, target_shape, do_ensemble=do_ensemble)
         out = Image.fromarray(np.asarray(out))
 
         nearest = image_in.resize(out.size, Image.NEAREST)
@@ -97,6 +120,7 @@ class TheraApp(DualVisionApp):
         }
         out_settings = {
             'scale': scale,
+            'model': model,
             'do_ensemble': do_ensemble,
         }
         return out_modalities, out_settings
@@ -147,6 +171,7 @@ class TheraApp(DualVisionApp):
             )
         if any(k not in results_settings for k in self.input_keys):
             raise gr.Error(f"Mismatching setgings keys")
+
         results_settings = {
             k: cls(**ctor_args, value=results_settings[k])
             for k, cls, ctor_args in zip(

gradio_dualvision/app_template.py

@@ -228,7 +228,7 @@ class DualVisionApp(gr.Blocks):
             }}
             #settings-accordion {{
                 margin: 0 auto;
-                max-width: 500px;
+                max-width: 650px;
             }}
         """
         if squeeze_canvas: