postprocessing_gfpgan.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536
  1. from PIL import Image
  2. import numpy as np
  3. from modules import scripts_postprocessing, gfpgan_model, ui_components
  4. import gradio as gr
  5. class ScriptPostprocessingGfpGan(scripts_postprocessing.ScriptPostprocessing):
  6. name = "GFPGAN"
  7. order = 2000
  8. def ui(self):
  9. with ui_components.InputAccordion(False, label="GFPGAN") as enable:
  10. gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Visibility", value=1.0, elem_id=self.elem_id_suffix("extras_gfpgan_visibility"))
  11. return {
  12. "enable": enable,
  13. "gfpgan_visibility": gfpgan_visibility,
  14. }
  15. def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, gfpgan_visibility):
  16. if gfpgan_visibility == 0 or not enable:
  17. return
  18. restored_img = gfpgan_model.gfpgan_fix_faces(np.array(pp.image.convert("RGB"), dtype=np.uint8))
  19. res = Image.fromarray(restored_img)
  20. if gfpgan_visibility < 1.0:
  21. if pp.image.size != res.size:
  22. res = res.resize(pp.image.size)
  23. if pp.image.mode != res.mode:
  24. res = res.convert(pp.image.mode)
  25. res = Image.blend(pp.image, res, gfpgan_visibility)
  26. pp.image = res
  27. pp.info["GFPGAN visibility"] = round(gfpgan_visibility, 3)