Add phone model (beta), allow models to use different architectures

app.py

@@ -18,7 +18,8 @@ with gr.Blocks(theme=gr.themes.Default().set(body_background_fill="#CCEEFF")) as
         modelSelect = gr.Dropdown(
             [
                 ["FM Radio Super Resolution","FM_Radio_SR.th"],
-                ["AM Radio Super Resolution (Beta)","AM_Radio_SR.th"] 
+                ["AM Radio Super Resolution (Beta)","AM_Radio_SR.th"],
+                ["Telephone Super Resolution (Beta)","Telephone_SR.th"] 
             ],
             label="Select Model:",
             value="FM_Radio_SR.th",
@@ -66,9 +67,19 @@ with gr.Blocks(theme=gr.themes.Default().set(body_background_fill="#CCEEFF")) as
             lrAudio = torch.tensor(audioData[1].copy().astype(np.float32)/32768).transpose(0,1)
         if audioData[0] != 44100:
             lrAudio = resample(lrAudio, audioData[0], 44100)
-        hrAudio=upscaleAudio(lrAudio, "models/" + model)
+        model_name, experiment_file = getModelInfo(model)
+        hrAudio=upscaleAudio(lrAudio, model, model_name=model_name, experiment_file=experiment_file)
         hrAudio=hrAudio / max(hrAudio.abs().max().item(), 1)
         outAudio=(hrAudio*32767).numpy().astype(np.int16).transpose(1,0)
         return tuple([44100, outAudio])
+    
+    def getModelInfo(modelFilename: str):
+        if(modelFilename == "FM_Radio_SR.th"):
+            return "aero", "aero_441-441_512_256.yaml"
+        if(modelFilename == "AM_Radio_SR.th"):
+            return "aero", "aero_441-441_512_256.yaml"
+        if(modelFilename == "Telephone_SR.th"):
+            return "aero", "aero_441-441_512_256.yaml"
+        return "aero", "aero_441-441_512_256.yaml"
 
 layout.launch()

models/Telephone_SR.th

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b59e32fccaf83c7e038b8c5e894eeebbce272e9ab00db6b20d45b2fba6e911ca
+size 136533968

processAudio.py

@@ -20,9 +20,9 @@ SEGMENT_DURATION_SEC = 5
 SEGMENT_OVERLAP_RATIO = 0.25
 SERIALIZE_KEY_STATE = 'state'
 
-def _load_model(checkpoint_file="models/FM_Radio_SR.th",model_name="aero"):
-    checkpoint_file = Path(checkpoint_file)
-    model = modelFactory.get_model(model_name)['generator']
+def _load_model(checkpoint_filename="FM_Radio_SR.th",model_name="aero",experiment_file="aero_441-441_512_256.yaml"):
+    checkpoint_file = Path("models/" + checkpoint_filename)
+    model = modelFactory.get_model(model_name,experiment_file)['generator']
     package = torch.load(checkpoint_file, 'cpu')
     if  'state' in package.keys(): #raw model file
         logger.info(bold(f'Loading model {model_name} from file.'))
@@ -35,9 +35,9 @@ def crossfade_and_blend(out_clip, in_clip, segment_overlap_samples):
     fade_in = torchaudio.transforms.Fade(segment_overlap_samples, 0)
     return fade_out(out_clip) + fade_in(in_clip)
 
-def upscaleAudio(lr_sig, checkpoint_file: str, sr=44100, hr_sr=44100, model_name="aero", progress=Progress()):
+def upscaleAudio(lr_sig, checkpoint_file: str, sr=44100, hr_sr=44100, model_name="aero", experiment_file="aero_441-441_512_256.yaml", progress=Progress()):
 
-    model = _load_model(checkpoint_file,model_name)
+    model = _load_model(checkpoint_file,model_name,experiment_file)
     device = torch.device('cpu')
     if torch.cuda.is_available():
         device = torch.device('cuda')

src/models/modelFactory.py

@@ -2,12 +2,12 @@ from src.models.aero import Aero
 from src.models.seanet import Seanet
 from yaml import safe_load
 
-def get_model(model_name="aero"):
+def get_model(model_name="aero", experiment_file="aero_441-441_512_256.yaml"):
     if model_name == 'aero':
-        with open("conf/experiment/aero_441-441_512_256.yaml") as f:
+        with open("conf/experiment/" + experiment_file) as f:
             generator = Aero(**safe_load(f)["aero"])
     elif model_name == 'seanet':
-        with open("conf/experiment/seanet_441-441.yaml") as f:
+        with open("conf/experiment/" + experiment_file) as f:
             generator = Seanet(**safe_load(f)["seanet"])
 
     models = {'generator': generator}