change examples

app.py

@@ -20,15 +20,17 @@ from sampler import DifIRSampler
 from ResizeRight.resize_right import resize
 from basicsr.utils.download_util import load_file_from_url
 
-def predict(im_path, background_enhance, face_upsample, upscale, started_timesteps):
-    cfg_path = 'configs/sample/iddpm_ffhq512_swinir.yaml'
+# setting configurations
+cfg_path = 'configs/sample/iddpm_ffhq512_swinir.yaml'
+configs = OmegaConf.load(cfg_path)
+configs.aligned = False
+
+# build the sampler for diffusion
+sampler_dist = DifIRSampler(configs)
 
-    # setting configurations
-    configs = OmegaConf.load(cfg_path)
-    configs.aligned = False
+def predict(im_path, background_enhance, face_upsample, upscale, started_timesteps):
     configs.background_enhance = background_enhance
     configs.face_upsample = face_upsample
-
     started_timesteps = int(started_timesteps)
     assert started_timesteps < int(configs.diffusion.params.timestep_respacing)
 
@@ -56,9 +58,6 @@ def predict(im_path, background_enhance, face_upsample, upscale, started_timeste
         upscale = 2  # avoid momory exceeded due to too large img resolution
     configs.detection.upscale = int(upscale)
 
-    # build the sampler for diffusion
-    sampler_dist = DifIRSampler(configs)
-
     image_restored, face_restored, face_cropped = sampler_dist.sample_func_bfr_unaligned(
             y0=im_lq,
             start_timesteps=started_timesteps,
@@ -71,7 +70,7 @@ def predict(im_path, background_enhance, face_upsample, upscale, started_timeste
         restored_image_dir.mkdir()
     # save the whole image
     save_path = restored_image_dir / Path(im_path).name
-    util_image.imwrite(image_restored, save_path, chn='bgr', dtype_in='uint8')
+    util_image.imwrite(image_restored, save_path, chn='rgb', dtype_in='uint8')
 
     return image_restored, str(save_path)
 
@@ -114,10 +113,10 @@ If you have any questions, please feel free to contact me via <b>[email protected]
 demo = gr.Interface(
     predict,
     inputs=[
-        gr.inputs.Image(type="filepath", label="Input"),
-        gr.inputs.Checkbox(default=True, label="Background_Enhance"),
-        gr.inputs.Checkbox(default=True, label="Face_Upsample"),
-        gr.inputs.Number(default=2, label="Rescaling_Factor (up to 4)"),
+        gr.Image(type="filepath", label="Input"),
+        gr.Checkbox(default=True, label="Background_Enhance"),
+        gr.Checkbox(default=True, label="Face_Upsample"),
+        gr.Number(default=2, label="Rescaling_Factor (up to 4)"),
         gr.Slider(1, 200, value=100, step=10, label='Realism-Fidelity Trade-off')
     ],
     outputs=[

sampler.py

@@ -54,7 +54,12 @@ class BaseSampler:
         torch.cuda.manual_seed_all(seed)
 
     def setup_dist(self):
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        if torch.cuda.is_available():
+            self.device = torch.device('cuda')
+            print(f'Runing on GPU...')
+        else:
+            self.device = torch.device('cpu')
+            print(f'Runing on CPU...')
         self.rank = 0
 
     def build_model(self):