add resize for arbitraty size

sampler.py

@@ -166,6 +166,11 @@ class DifIRSampler(BaseSampler):
         # basical image restoration
         device = next(self.model.parameters()).device
         y0 = y0.to(device=device, dtype=torch.float32)
+
+        h_old, w_old = y0.shape[2:4]
+        if not (h_old == self.configs.im_size and w_old == self.configs.im_size):
+            y0 = resize(y0, out_shape=(self.configs.im_size,) * 2).to(torch.float32)
+
         if need_restoration:
             with torch.no_grad():
                 if model_kwargs_ir is None:
@@ -176,10 +181,6 @@ class DifIRSampler(BaseSampler):
             im_hq = y0
         im_hq.clamp_(0.0, 1.0)
 
-        h_old, w_old = im_hq.shape[2:4]
-        if not (h_old == self.configs.im_size and w_old == self.configs.im_size):
-            im_hq = resize(im_hq, out_shape=(self.configs.im_size,) * 2).to(torch.float32)
-
         # diffuse for im_hq
         yt = self.diffusion.q_sample(
                 x_start=post_fun(im_hq),