Spaces:
Runtime error
Runtime error
File size: 2,694 Bytes
3712f32 |
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 |
import os
from PIL import Image
import numpy as np
import roop.globals
from roop.core import start, decode_execution_providers, suggest_max_memory, suggest_execution_threads
from roop.processors.frame.core import get_frame_processors_modules
def to_pil(img):
"""Convert np.ndarray to PIL.Image if needed."""
if isinstance(img, np.ndarray):
return Image.fromarray(img)
return img
def swap_face_now(source_img, target_img, do_enhance=True):
TEMP_DIR = "temp/"
os.makedirs(TEMP_DIR, exist_ok=True)
# ✅ Ép kiểu an toàn và convert về RGB để tránh lỗi JPEG
source_img = to_pil(source_img)
target_img = to_pil(target_img)
# Convert về RGB nếu ảnh có alpha channel
if source_img.mode in ('RGBA', 'LA'):
# Tạo background trắng và paste ảnh lên
rgb_source = Image.new('RGB', source_img.size, (255, 255, 255))
rgb_source.paste(source_img, mask=source_img.split()[-1] if source_img.mode == 'RGBA' else None)
source_img = rgb_source
elif source_img.mode != 'RGB':
source_img = source_img.convert('RGB')
if target_img.mode in ('RGBA', 'LA'):
# Tạo background trắng và paste ảnh lên
rgb_target = Image.new('RGB', target_img.size, (255, 255, 255))
rgb_target.paste(target_img, mask=target_img.split()[-1] if target_img.mode == 'RGBA' else None)
target_img = rgb_target
elif target_img.mode != 'RGB':
target_img = target_img.convert('RGB')
source_path = os.path.join(TEMP_DIR, "input.jpg")
target_path = os.path.join(TEMP_DIR, "target.jpg")
output_path = os.path.join(TEMP_DIR, "output.jpg")
source_img.save(source_path, 'JPEG')
target_img.save(target_path, 'JPEG')
# Roop config
roop.globals.source_path = source_path
roop.globals.target_path = target_path
roop.globals.output_path = output_path
roop.globals.frame_processors = ["face_swapper", "face_enhancer"] if do_enhance else ["face_swapper"]
roop.globals.headless = True
roop.globals.many_faces = False
roop.globals.max_memory = suggest_max_memory()
roop.globals.execution_providers = decode_execution_providers(["cuda"])
roop.globals.execution_threads = suggest_execution_threads()
roop.globals.keep_fps = False
roop.globals.keep_audio = False
roop.globals.keep_frames = False
roop.globals.video_encoder = None
for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
if not frame_processor.pre_check():
return None
start()
if not os.path.exists(output_path):
return None
return Image.open(output_path)
|