postprocessing_gfpgan.py 1.0 KB

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