add alinged case

app.py

@@ -62,79 +62,92 @@ 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)
 
-    image_restored, face_restored, face_cropped = sampler_dist.sample_func_bfr_unaligned(
-            y0=im_lq,
-            start_timesteps=started_timesteps,
-            need_restoration=True,
-            draw_box=False,
-            )
+    if background_enhance:
+        image_restored, face_restored, face_cropped = sampler_dist.sample_func_bfr_unaligned(
+                y0=im_lq,
+                start_timesteps=started_timesteps,
+                need_restoration=True,
+                draw_box=False,
+                )   # h x w x c, numpy array, [0, 255], uint8, BGR
+        image_restored = util_image.bgr2rgb(image_restored)
+    else:
+        image_restored = sampler_dist.sample_func_ir_aligned(
+                y0=im_lq,
+                start_timesteps=started_timesteps,
+                need_restoration=True,
+                )[0]  # b x c x h x w, [0, 1], torch tensor, RGB
+        image_restored = util_image.tensor2img(
+                image_restored.cpu(),
+                rgb2bgr=False,
+                out_type=np.uint8,
+                min_max=(0, 1),
+                )     # h x w x c, [0, 255], uint8, RGB, numpy array
 
     restored_image_dir = Path('restored_output')
     if not restored_image_dir.exists():
         restored_image_dir.mkdir()
     # save the whole image
     save_path = restored_image_dir / Path(im_path).name
-    image_restored = util_image.bgr2rgb(image_restored)
     util_image.imwrite(image_restored, save_path, chn='rgb', dtype_in='uint8')
 
     return image_restored, str(save_path)
 
-title = "DifFace: Blind Face Restoration with Diffused Error Contraction"
-description = r"""
-<b>Official Gradio demo</b> for <a href='https://github.com/zsyOAOA/DifFace' target='_blank'><b>DifFace: Blind Face Restoration with Diffused Error Contraction</b></a>.<br>
-🔥 DifFace is a robust face restoration algorithm for old or corrupted photos.<br>
-"""
-article = r"""
-If DifFace is helpful for your work, please help to ⭐ the <a href='https://github.com/zsyOAOA/DifFace' target='_blank'>Github Repo</a>. Thanks!
-[![GitHub Stars](https://img.shields.io/github/stars/zsyOAOA/DifFace?affiliations=OWNER&color=green&style=social)](https://github.com/zsyOAOA/DifFace)
-
----
-
-📝 **Citation**
-
-If our work is useful for your research, please consider citing:
-```bibtex
-@article{yue2022difface,
-  title={DifFace: Blind Face Restoration with Diffused Error Contraction},
-  author={Yue, Zongsheng and Loy, Chen Change},
-  journal={arXiv preprint arXiv:2212.06512},
-  year={2022}
-}
-```
-
-📋 **License**
-
-This project is licensed under <a rel="license" href="https://github.com/zsyOAOA/DifFace/blob/master/LICENSE">S-Lab License 1.0</a>.
-Redistribution and use for non-commercial purposes should follow this license.
-
-📧 **Contact**
-If you have any questions, please feel free to contact me via <b>zsyzam@gmail.com</b>.
-![visitors](https://visitor-badge.laobi.icu/badge?page_id=zsyOAOA/DifFace)
-"""
-
-demo = gr.Interface(
-    predict,
-    inputs=[
-        gr.Image(type="filepath", label="Input"),
-        gr.Checkbox(value=True, label="Background_Enhance"),
-        gr.Checkbox(value=True, label="Face_Upsample"),
-        gr.Number(value=2, label="Rescaling_Factor (up to 4)"),
-        gr.Slider(1, 160, value=80, step=10, label='Realism-Fidelity Trade-off (1 for better fidelity, 160 for better realism)')
-    ],
-    outputs=[
-        gr.Image(type="numpy", label="Output"),
-        gr.outputs.File(label="Download the output")
-    ],
-    title=title,
-    description=description,
-    article=article,
-    examples=[
-        ['./testdata/whole_imgs/00.jpg', True, True, 2, 80],
-        ['./testdata/whole_imgs/01.jpg', True, True, 2, 80],
-        ['./testdata/whole_imgs/04.jpg', True, True, 2, 80],
-        ['./testdata/whole_imgs/05.jpg', True, True, 2, 80],
-      ]
-    )
-
-demo.queue(concurrency_count=4)
-demo.launch()
+# title = "DifFace: Blind Face Restoration with Diffused Error Contraction"
+# description = r"""
+# <b>Official Gradio demo</b> for <a href='https://github.com/zsyOAOA/DifFace' target='_blank'><b>DifFace: Blind Face Restoration with Diffused Error Contraction</b></a>.<br>
+# 🔥 DifFace is a robust face restoration algorithm for old or corrupted photos.<br>
+# """
+# article = r"""
+# If DifFace is helpful for your work, please help to ⭐ the <a href='https://github.com/zsyOAOA/DifFace' target='_blank'>Github Repo</a>. Thanks!
+# [![GitHub Stars](https://img.shields.io/github/stars/zsyOAOA/DifFace?affiliations=OWNER&color=green&style=social)](https://github.com/zsyOAOA/DifFace)
+
+# ---
+
+# 📝 **Citation**
+
+# If our work is useful for your research, please consider citing:
+# ```bibtex
+# @article{yue2022difface,
+  # title={DifFace: Blind Face Restoration with Diffused Error Contraction},
+  # author={Yue, Zongsheng and Loy, Chen Change},
+  # journal={arXiv preprint arXiv:2212.06512},
+  # year={2022}
+# }
+# ```
+
+# 📋 **License**
+
+# This project is licensed under <a rel="license" href="https://github.com/zsyOAOA/DifFace/blob/master/LICENSE">S-Lab License 1.0</a>.
+# Redistribution and use for non-commercial purposes should follow this license.
+
+# 📧 **Contact**
+# If you have any questions, please feel free to contact me via <b>zsyzam@gmail.com</b>.
+# ![visitors](https://visitor-badge.laobi.icu/badge?page_id=zsyOAOA/DifFace)
+# """
+
+# demo = gr.Interface(
+    # predict,
+    # inputs=[
+        # gr.Image(type="filepath", label="Input"),
+        # gr.Checkbox(value=True, label="Background_Enhance"),
+        # gr.Checkbox(value=True, label="Face_Upsample"),
+        # gr.Number(value=2, label="Rescaling_Factor (up to 4)"),
+        # gr.Slider(1, 160, value=80, step=10, label='Realism-Fidelity Trade-off')
+    # ],
+    # outputs=[
+        # gr.Image(type="numpy", label="Output"),
+        # gr.outputs.File(label="Download the output")
+    # ],
+    # title=title,
+    # description=description,
+    # article=article,
+    # examples=[
+        # ['./testdata/whole_imgs/00.jpg', True, True, 2, 80],
+        # ['./testdata/whole_imgs/01.jpg', True, True, 2, 80],
+        # ['./testdata/whole_imgs/04.jpg', True, True, 2, 80],
+        # ['./testdata/whole_imgs/05.jpg', True, True, 2, 80],
+      # ]
+    # )
+
+# demo.queue(concurrency_count=4)
+# demo.launch()