with torch.no_grad():

app.py

@@ -386,7 +386,8 @@ def restore_on_gpu(
     LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
     captions = ['']
 
-    samples = model.batchify_sample(LQ, captions, num_steps=edm_steps, restoration_scale=s_stage1, s_churn=s_churn,
+    with torch.no_grad():
+        samples = model.batchify_sample(LQ, captions, num_steps=edm_steps, restoration_scale=s_stage1, s_churn=s_churn,
                                     s_noise=s_noise, cfg_scale=s_cfg, control_scale=s_stage2, seed=seed,
                                     num_samples=num_samples, p_p=a_prompt, n_p=n_prompt, color_fix_type=color_fix_type,
                                     use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,