File size: 4,913 Bytes
55f6076 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import os, glob
from PIL import Image
import modules.scripts as scripts
# from modules.upscaler import Upscaler, UpscalerData
from modules import scripts, scripts_postprocessing
from modules.processing import (
StableDiffusionProcessing,
StableDiffusionProcessingImg2Img,
)
from modules.shared import state
from scripts.reactor_logger import logger
from scripts.reactor_swapper import swap_face, get_current_faces_model, analyze_faces, half_det_size
import folder_paths
import comfy.model_management as model_management
def get_models():
models_path = os.path.join(folder_paths.models_dir,"insightface/*")
models = glob.glob(models_path)
models = [x for x in models if x.endswith(".onnx") or x.endswith(".pth")]
return models
class FaceSwapScript(scripts.Script):
def process(
self,
p: StableDiffusionProcessing,
img,
enable,
source_faces_index,
faces_index,
model,
swap_in_source,
swap_in_generated,
gender_source,
gender_target,
face_model,
):
self.enable = enable
if self.enable:
self.source = img
self.swap_in_generated = swap_in_generated
self.gender_source = gender_source
self.gender_target = gender_target
self.model = model
self.face_model = face_model
self.source_faces_index = [
int(x) for x in source_faces_index.strip(",").split(",") if x.isnumeric()
]
self.faces_index = [
int(x) for x in faces_index.strip(",").split(",") if x.isnumeric()
]
if len(self.source_faces_index) == 0:
self.source_faces_index = [0]
if len(self.faces_index) == 0:
self.faces_index = [0]
if self.gender_source is None or self.gender_source == "no":
self.gender_source = 0
elif self.gender_source == "female":
self.gender_source = 1
elif self.gender_source == "male":
self.gender_source = 2
if self.gender_target is None or self.gender_target == "no":
self.gender_target = 0
elif self.gender_target == "female":
self.gender_target = 1
elif self.gender_target == "male":
self.gender_target = 2
# if self.source is not None:
if isinstance(p, StableDiffusionProcessingImg2Img) and swap_in_source:
logger.status(f"Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
for i in range(len(p.init_images)):
if state.interrupted or model_management.processing_interrupted():
logger.status("Interrupted by User")
break
if len(p.init_images) > 1:
logger.status(f"Swap in %s", i)
result = swap_face(
self.source,
p.init_images[i],
source_faces_index=self.source_faces_index,
faces_index=self.faces_index,
model=self.model,
gender_source=self.gender_source,
gender_target=self.gender_target,
face_model=self.face_model,
)
p.init_images[i] = result
logger.status("--Done!--")
# else:
# logger.error(f"Please provide a source face")
def postprocess_batch(self, p, *args, **kwargs):
if self.enable:
images = kwargs["images"]
def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args):
if self.enable and self.swap_in_generated:
if self.source is not None:
logger.status(f"Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
image: Image.Image = script_pp.image
result = swap_face(
self.source,
image,
source_faces_index=self.source_faces_index,
faces_index=self.faces_index,
model=self.model,
upscale_options=self.upscale_options,
gender_source=self.gender_source,
gender_target=self.gender_target,
)
try:
pp = scripts_postprocessing.PostprocessedImage(result)
pp.info = {}
p.extra_generation_params.update(pp.info)
script_pp.image = pp.image
except:
logger.error(f"Cannot create a result image")
|