402 lines
13 KiB
Python
402 lines
13 KiB
Python
# roop_custom_nodes.py
|
|
import os
|
|
import subprocess
|
|
import glob
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
import hmac
|
|
import hashlib
|
|
import requests
|
|
|
|
# Helper: save tensor image
|
|
def save_tensor_as_image(tensor, path):
|
|
np_img = (tensor[0].cpu().numpy() * 255).astype(np.uint8)
|
|
if np_img.shape[0] == 3: # If format is CHW
|
|
np_img = np.transpose(np_img, (1, 2, 0)) # Convert to HWC for PIL
|
|
img = Image.fromarray(np_img)
|
|
img.save(path)
|
|
|
|
# Helper: load image as tensor
|
|
def load_image_as_tensor(path):
|
|
img = Image.open(path).convert("RGB")
|
|
np_img = np.array(img).astype(np.float32) / 255.0
|
|
return torch.from_numpy(np_img).unsqueeze(0)
|
|
|
|
def get_unique_filename(path):
|
|
base, ext = os.path.splitext(path)
|
|
counter = 1
|
|
new_path = path
|
|
while os.path.exists(new_path):
|
|
new_path = f"{base}_{counter}{ext}"
|
|
counter += 1
|
|
return new_path
|
|
|
|
def send_webhook_image(webhook_url, webhook_secret, output_path, extra_data=None):
|
|
if not webhook_url:
|
|
return
|
|
|
|
with open(output_path, 'rb') as f:
|
|
image_data = f.read()
|
|
|
|
# GitHub-style HMAC signature
|
|
signature = 'sha256=' + hmac.new(
|
|
webhook_secret.encode('utf-8'),
|
|
image_data,
|
|
hashlib.sha256
|
|
).hexdigest()
|
|
|
|
headers = {
|
|
'X-Hub-Signature-256': signature
|
|
}
|
|
|
|
files = {
|
|
'file': ('swapped.png', image_data, 'image/png')
|
|
}
|
|
|
|
data = extra_data or {}
|
|
|
|
try:
|
|
resp = requests.post(webhook_url, headers=headers, files=files, data=data)
|
|
resp.raise_for_status()
|
|
print(f"[Webhook] Sent successfully to {webhook_url}")
|
|
except Exception as e:
|
|
print(f"[Webhook Error] {e}")
|
|
|
|
|
|
class RoopFaceSwap:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source_image": ("IMAGE",),
|
|
"target_image": ("IMAGE",),
|
|
"roop_dir": ("STRING", {"default": "/content/roop"}),
|
|
"output_name": ("STRING", {"default": "roop_output.png"}),
|
|
"many_faces": ("BOOLEAN", {"default": False})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("swapped_image",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "Roop/Basic"
|
|
|
|
def run(self, source_image, target_image, roop_dir, output_name, many_faces):
|
|
temp_dir = os.path.join(roop_dir, "temp_io")
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
|
|
source_path = os.path.join(temp_dir, "source.png")
|
|
target_path = os.path.join(temp_dir, "target.png")
|
|
output_path = get_unique_filename(os.path.join(temp_dir, output_name))
|
|
|
|
save_tensor_as_image(source_image, source_path)
|
|
save_tensor_as_image(target_image, target_path)
|
|
|
|
cmd = [
|
|
"python", "run.py",
|
|
"-s", source_path,
|
|
"-t", target_path,
|
|
"-o", output_path,
|
|
"--execution-provider", "cuda",
|
|
"--frame-processor", "face_swapper"
|
|
]
|
|
if many_faces:
|
|
cmd.append("--many-faces")
|
|
|
|
subprocess.run(cmd, check=True, cwd=roop_dir)
|
|
|
|
if not os.path.exists(output_path):
|
|
print(f"[Warning] Roop did not produce output: {output_path}")
|
|
blank = torch.zeros_like(target_image)
|
|
return (blank,)
|
|
|
|
return (load_image_as_tensor(output_path),)
|
|
|
|
|
|
class RoopFaceSwapWithEnhancer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source_image": ("IMAGE",),
|
|
"target_image": ("IMAGE",),
|
|
"roop_dir": ("STRING", {"default": "/content/roop"}),
|
|
"output_name": ("STRING", {"default": "roop_output.png"}),
|
|
"many_faces": ("BOOLEAN", {"default": False})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("swapped_image",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "Roop/Enhanced"
|
|
|
|
def run(self, source_image, target_image, roop_dir, output_name, many_faces):
|
|
temp_dir = os.path.join(roop_dir, "temp_io")
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
|
|
source_path = os.path.join(temp_dir, "source.png")
|
|
target_path = os.path.join(temp_dir, "target.png")
|
|
output_path = get_unique_filename(os.path.join(temp_dir, output_name))
|
|
|
|
save_tensor_as_image(source_image, source_path)
|
|
save_tensor_as_image(target_image, target_path)
|
|
|
|
cmd = [
|
|
"python", "run.py",
|
|
"-s", source_path,
|
|
"-t", target_path,
|
|
"-o", output_path,
|
|
"--execution-provider", "cuda",
|
|
"--frame-processor", "face_swapper", "face_enhancer"
|
|
]
|
|
if many_faces:
|
|
cmd.append("--many-faces")
|
|
|
|
subprocess.run(cmd, check=True, cwd=roop_dir)
|
|
|
|
if not os.path.exists(output_path):
|
|
print(f"[Warning] Roop did not produce output: {output_path}")
|
|
blank = torch.zeros_like(target_image)
|
|
return (blank,)
|
|
|
|
return (load_image_as_tensor(output_path),)
|
|
|
|
|
|
class RoopBatchFaceSwap:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source_image": ("IMAGE",),
|
|
"input_dir": ("STRING", {"default": "/input/images"}),
|
|
"output_dir": ("STRING", {"default": "/output/images"}),
|
|
"roop_dir": ("STRING", {"default": "/content/roop"}),
|
|
"use_enhancer": ("BOOLEAN", {"default": False}),
|
|
"many_faces": ("BOOLEAN", {"default": False})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
RETURN_NAMES = ()
|
|
FUNCTION = "run"
|
|
CATEGORY = "Roop/Batch"
|
|
|
|
def run(self, source_image, input_dir, output_dir, roop_dir, use_enhancer, many_faces):
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
temp_dir = os.path.join(roop_dir, "temp_io")
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
|
|
source_path = os.path.join(temp_dir, "source.png")
|
|
save_tensor_as_image(source_image, source_path)
|
|
|
|
image_paths = glob.glob(os.path.join(input_dir, "*.jpg")) + \
|
|
glob.glob(os.path.join(input_dir, "*.png"))
|
|
|
|
for img_path in image_paths:
|
|
target_name = os.path.basename(img_path)
|
|
target_path = os.path.join(temp_dir, "target.png")
|
|
output_path = os.path.join(output_dir, f"out_{target_name}")
|
|
|
|
Image.open(img_path).save(target_path)
|
|
|
|
cmd = [
|
|
"python", "run.py",
|
|
"-s", source_path,
|
|
"-t", target_path,
|
|
"-o", output_path,
|
|
"--execution-provider", "cuda",
|
|
"--frame-processor", "face_swapper"
|
|
]
|
|
|
|
if use_enhancer:
|
|
cmd[-1] += " face_enhancer"
|
|
if many_faces:
|
|
cmd.append("--many-faces")
|
|
|
|
subprocess.run(cmd, check=True, cwd=roop_dir)
|
|
|
|
if not os.path.exists(output_path):
|
|
print(f"[Skipped] NSFW or error: {img_path} -> No output generated.")
|
|
continue
|
|
|
|
return ()
|
|
|
|
|
|
class RoopFaceSwapVideo:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source_image": ("IMAGE",),
|
|
"target_video_path": ("STRING", {"default": "/path/to/video.mp4"}),
|
|
"roop_dir": ("STRING", {"default": "/content/roop"}),
|
|
"output_name": ("STRING", {"default": "swapped_video.mp4"}),
|
|
"many_faces": ("BOOLEAN", {"default": False})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
RETURN_NAMES = ()
|
|
FUNCTION = "run"
|
|
CATEGORY = "Roop/Video"
|
|
|
|
def run(self, source_image, target_video_path, roop_dir, output_name, many_faces):
|
|
temp_dir = os.path.join(roop_dir, "temp_io")
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
|
|
source_path = os.path.join(temp_dir, "source.png")
|
|
output_path = get_unique_filename(os.path.join(temp_dir, output_name))
|
|
save_tensor_as_image(source_image, source_path)
|
|
|
|
cmd = [
|
|
"python", "run.py",
|
|
"-s", source_path,
|
|
"-t", target_video_path,
|
|
"-o", output_path,
|
|
"--keep-fps",
|
|
"--keep-frames",
|
|
"--execution-provider", "cuda",
|
|
"--frame-processor", "face_swapper"
|
|
]
|
|
if many_faces:
|
|
cmd.append("--many-faces")
|
|
|
|
subprocess.run(cmd, check=True, cwd=roop_dir)
|
|
|
|
if not os.path.exists(output_path):
|
|
print(f"[Warning] Roop did not produce video output: {output_path}")
|
|
|
|
return ()
|
|
|
|
|
|
class RoopSendWebhookImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_tensor": ("IMAGE",),
|
|
"filename": ("STRING", {"default": "output.png"}),
|
|
"webhook_url": ("STRING", {"default": ""}),
|
|
"webhook_secret": ("STRING", {"default": ""}),
|
|
"enable_webhook": ("BOOLEAN", {"default": True}),
|
|
"roop_dir": ("STRING", {"default": "/content/roop"})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
RETURN_NAMES = ()
|
|
FUNCTION = "run"
|
|
CATEGORY = "Roop/Webhook"
|
|
|
|
def run(self, image_tensor, filename, webhook_url, webhook_secret, enable_webhook, roop_dir):
|
|
if not enable_webhook or not webhook_url:
|
|
print("[WebhookImage] Disabled or URL not set — skipping.")
|
|
return ()
|
|
|
|
temp_dir = os.path.join(roop_dir, "temp_io")
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
output_path = os.path.join(temp_dir, filename)
|
|
|
|
save_tensor_as_image(image_tensor, output_path)
|
|
|
|
try:
|
|
with open(output_path, 'rb') as f:
|
|
file_data = f.read()
|
|
|
|
headers = {}
|
|
if webhook_secret:
|
|
signature = 'sha256=' + hmac.new(
|
|
webhook_secret.encode('utf-8'),
|
|
file_data,
|
|
hashlib.sha256
|
|
).hexdigest()
|
|
headers['X-Hub-Signature-256'] = signature
|
|
|
|
files = {
|
|
'file': (filename, file_data, 'image/png')
|
|
}
|
|
|
|
resp = requests.post(webhook_url, headers=headers, files=files)
|
|
resp.raise_for_status()
|
|
print(f"[WebhookImage] Sent image: {filename} → {webhook_url}")
|
|
except Exception as e:
|
|
print(f"[WebhookImage Error] {e}")
|
|
|
|
return ()
|
|
|
|
|
|
class RoopSendWebhookFile:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"file_path": ("STRING",),
|
|
"filename": ("STRING", {"default": "output.mp4"}),
|
|
"webhook_url": ("STRING", {"default": ""}),
|
|
"webhook_secret": ("STRING", {"default": ""}),
|
|
"enable_webhook": ("BOOLEAN", {"default": True})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
RETURN_NAMES = ()
|
|
FUNCTION = "run"
|
|
CATEGORY = "Roop/Webhook"
|
|
|
|
def run(self, file_path, filename, webhook_url, webhook_secret, enable_webhook):
|
|
if not enable_webhook or not webhook_url:
|
|
print("[WebhookFile] Disabled or URL not set — skipping.")
|
|
return ()
|
|
|
|
if not os.path.exists(file_path):
|
|
print(f"[WebhookFile] File does not exist: {file_path}")
|
|
return ()
|
|
|
|
try:
|
|
with open(file_path, 'rb') as f:
|
|
file_data = f.read()
|
|
|
|
headers = {}
|
|
if webhook_secret:
|
|
signature = 'sha256=' + hmac.new(
|
|
webhook_secret.encode('utf-8'),
|
|
file_data,
|
|
hashlib.sha256
|
|
).hexdigest()
|
|
headers['X-Hub-Signature-256'] = signature
|
|
|
|
content_type = "video/mp4" if filename.endswith(".mp4") else "application/octet-stream"
|
|
files = {
|
|
'file': (filename, file_data, content_type)
|
|
}
|
|
|
|
resp = requests.post(webhook_url, headers=headers, files=files)
|
|
resp.raise_for_status()
|
|
print(f"[WebhookFile] Sent file: {filename} → {webhook_url}")
|
|
except Exception as e:
|
|
print(f"[WebhookFile Error] {e}")
|
|
|
|
return ()
|
|
|
|
|
|
# Register with ComfyUI
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RoopFaceSwap": RoopFaceSwap,
|
|
"RoopFaceSwapWithEnhancer": RoopFaceSwapWithEnhancer,
|
|
"RoopBatchFaceSwap": RoopBatchFaceSwap,
|
|
"RoopFaceSwapVideo": RoopFaceSwapVideo,
|
|
"RoopSendWebhookImage": RoopSendWebhookImage,
|
|
"RoopSendWebhookFile": RoopSendWebhookFile,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RoopFaceSwap": "Roop Face Swap (Image)",
|
|
"RoopFaceSwapWithEnhancer": "Roop Face Swap + Enhancer",
|
|
"RoopBatchFaceSwap": "Roop Batch Image Folder",
|
|
"RoopFaceSwapVideo": "Roop Face Swap (Video)",
|
|
"RoopSendWebhookImage": "Roop Webhook: Image Tensor",
|
|
"RoopSendWebhookFile": "Roop Webhook: File Path",
|
|
}
|