Add EDSR model

.ci/docker/requirements-ci.txt

@@ -5,6 +5,7 @@ ruamel.yaml==0.17.32
 sympy==1.12
 timm==0.6.13
 tomli==2.0.1
+torchsr==1.0.4
 transformers==4.31.0
 zstd==1.5.5.1
 pytest==7.2.0

examples/models/TARGETS

@@ -10,6 +10,7 @@ python_library(
         "//caffe2:torch",
         "//executorch/examples/models:model_base",  # @manual
         "//executorch/examples/models/deeplab_v3:dl3_model",  # @manual
+        "//executorch/examples/models/edsr:edsr_model",  # @manual
         "//executorch/examples/models/emformer_rnnt:emformer_rnnt_model",  # @manual
         "//executorch/examples/models/inception_v3:ic3_model",  # @manual
         "//executorch/examples/models/inception_v4:ic4_model",  # @manual

examples/models/__init__.py

@@ -12,6 +12,7 @@
     "add": ("toy_model", "AddModule"),
     "add_mul": ("toy_model", "AddMulModule"),
     "dl3": ("deeplab_v3", "DeepLabV3ResNet50Model"),
+    "edsr": ("edsr", "EdsrModel"),
     "emformer_transcribe": ("emformer_rnnt", "EmformerRnntTranscriberModel"),
     "emformer_predict": ("emformer_rnnt", "EmformerRnntPredictorModel"),
     "emformer_join": ("emformer_rnnt", "EmformerRnntJoinerModel"),

examples/models/edsr/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .model import EdsrModel
+
+__all__ = [
+    EdsrModel,
+]

examples/models/edsr/model.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+
+import torch
+
+from torchsr.models import edsr_r16f64  # @manual
+
+from ..model_base import EagerModelBase
+
+
+class EdsrModel(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading edsr model")
+        m = edsr_r16f64(2, True)
+        logging.info("Loaded edsr model")
+        return m
+
+    def get_example_inputs(self):
+        tensor_size = (1, 3, 224, 224)
+        return (torch.randn(tensor_size),)

install_requirements.sh

@@ -33,3 +33,6 @@ pip install --pre timm==${TIMM_VERSION}
 
 TRANSFORMERS_VERSION=4.31.0
 pip install --pre transformers==${TRANSFORMERS_VERSION}
+
+TORCHSR_VERSION=1.0.4
+pip install --pre torchsr==${TORCHSR_VERSION}