test

app.py

@@ -31,7 +31,7 @@ def torch_to_pil(image):
 
 
 def image_mod(image, noise_level, denoiser):
-    image = pil_to_torch(image, ref_size=128 if denoiser == 'SwinIR' else 256)
+    image = pil_to_torch(image, ref_size=128 if denoiser == 'DiffUNet' else 256)
     if denoiser == 'DnCNN':
         den = dinv.models.DnCNN()
         sigma0 = 2/255
@@ -46,8 +46,8 @@ def image_mod(image, noise_level, denoiser):
         denoiser = dinv.models.TGVDenoiser()
     elif denoiser == 'Wavelets':
         denoiser = dinv.models.WaveletPrior()
-    elif denoiser == 'SwinIR':
-        denoiser = dinv.models.SwinIR()
+    elif denoiser == 'DiffUNet':
+        denoiser = dinv.models.DiffUNet()
     elif denoiser == 'DRUNet':
         denoiser = dinv.models.DRUNet()
     else:
@@ -64,7 +64,7 @@ input_image_output = gr.Image(label='Input Image')
 
 noise_levels = gr.Dropdown(choices=[0.1, 0.2, 0.3, 0.5, 1], value=0.1, label='Noise Level')
 
-denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'SwinIR', 'BM3D', 'MedianFilter', 'TV', 'TGV', 'Wavelets'], value='DnCNN', label='Denoiser')
+denoiser = gr.Dropdown(choices=['DnCNN', 'DRUNet', 'DiffUNet', 'BM3D', 'MedianFilter', 'TV', 'TGV', 'Wavelets'], value='DnCNN', label='Denoiser')
 
 demo = gr.Interface(
     image_mod,