import gradio as gr import spaces import os import cv2 import torch from PIL import Image from insightface.app import FaceAnalysis from ip_adapter.ip_adapter_faceid import IPAdapterFaceID from transformers import CLIPFeatureExtractor from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler, AutoencoderKL from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker import dlib import imutils from imutils import face_utils import numpy as np from skimage import transform as tf import random base_model_path = "SG161222/Realistic_Vision_V6.0_B1_noVAE" base_cache = "model-cache" vae_model_path = "stabilityai/sd-vae-ft-mse" ip_cache = "./ip-cache" device = "cuda" # Setup function to load models and other dependencies def setup(): """Load the model into memory to make running multiple predictions efficient""" # Get ip-adapter-faceid model if not os.path.exists("ip-cache/ip-adapter-faceid_sd15.bin"): os.makedirs(ip_cache, exist_ok=True) os.system(f"wget -O ip-cache/ip-adapter-faceid_sd15.bin https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid_sd15.bin") # Download shape_predictor_68_face_landmarks.dat if it doesn't exist if not os.path.exists("faceid/shape_predictor_68_face_landmarks.dat"): os.makedirs("faceid", exist_ok=True) os.system("wget -O faceid/shape_predictor_68_face_landmarks.dat https://github.com/italojs/facial-landmarks-recognition/raw/master/shape_predictor_68_face_landmarks.dat") # Face embedding app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) app.prepare(ctx_id=0, det_size=(640, 640)) # SD noise_scheduler = EulerDiscreteScheduler( num_train_timesteps=1000, beta_start=0.00085, beta_end=0.012 ) vae = AutoencoderKL.from_pretrained( vae_model_path ).to(dtype=torch.float16) pipe = StableDiffusionPipeline.from_pretrained( base_model_path, torch_dtype=torch.float16, scheduler=noise_scheduler, vae=vae, feature_extractor=CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32"), safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker"), cache_dir=base_cache, ) pipe = pipe.to(device) # IP adapter ip_model = IPAdapterFaceID( pipe, "ip-cache/ip-adapter-faceid_sd15.bin", device ) return app, ip_model app, ip_model = setup() def get_face_landmarks(image_path): def add_padding(image, padding_size=50): height, width = image.shape[:2] padded_image = cv2.copyMakeBorder( image, top=padding_size, bottom=padding_size, left=padding_size, right=padding_size, borderType=cv2.BORDER_CONSTANT, value=[255, 255, 255] # White padding ) return padded_image detector = dlib.get_frontal_face_detector() predictor = dlib.shape_predictor('faceid/shape_predictor_68_face_landmarks.dat') image = cv2.imread(image_path) image = imutils.resize(image, width=512) # Add padding to the image image = add_padding(image) gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) rects = detector(gray, 1) for (i, rect) in enumerate(rects): shape = predictor(gray, rect) shape = face_utils.shape_to_np(shape) return shape, image def morph_faces(image1_path, image2_path, alpha=0.5): landmarks1, image1 = get_face_landmarks(image1_path) landmarks2, image2 = get_face_landmarks(image2_path) average_landmarks = (landmarks1 + landmarks2) / 2 tform1 = tf.estimate_transform('similarity', landmarks1, average_landmarks) tform2 = tf.estimate_transform('similarity', landmarks2, average_landmarks) morphed_image1 = tf.warp(image1, inverse_map=tform1.inverse, output_shape=(512, 512)) morphed_image2 = tf.warp(image2, inverse_map=tform2.inverse, output_shape=(512, 512)) morphed_image = (1 - alpha) * morphed_image1 + alpha * morphed_image2 morphed_image = (morphed_image * 255).astype(np.uint8) # Convert to [0, 255] range output_path = "tmp.png" cv2.imwrite(output_path, morphed_image) return output_path def get_negative_prompt(gender): if gender == "Boy": return "(mascara, makeup: 1.4), (breasts, boobs, naked, nude: 1.4), (deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation" elif gender == "Girl": return "(beard, mustache, male features: 1.4), (naked, nude: 1.4), (deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation" else: # Random return "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation" def construct_prompt(base_prompt, additional_prompt, gender): full_prompt = base_prompt if gender == "Boy": full_prompt += ", male child, boy" elif gender == "Girl": full_prompt += ", female child, girl" if additional_prompt: full_prompt += ", " + additional_prompt return full_prompt @spaces.GPU(duration = 40) def generate_image(face_image_1, face_image_2, additional_prompt, gender): base_prompt = "portrait of a 6 y.o. child, 8k, HD, happy, perfect eyes, cute" full_prompt = construct_prompt(base_prompt, additional_prompt, gender) negative_prompt = get_negative_prompt(gender) baby_image_path = morph_faces(face_image_1, face_image_2) def generate_images(faceid_embeds, num_outputs=1): images = ip_model.generate( prompt=full_prompt, negative_prompt=negative_prompt, faceid_embeds=faceid_embeds, num_samples=num_outputs, width=768, height=768, num_inference_steps=40, seed=None ) return images faceid_embeds = app.get(cv2.imread(baby_image_path))[0].normed_embedding faceid_embeds = torch.from_numpy(faceid_embeds).unsqueeze(0) generated_images = generate_images(faceid_embeds) return generated_images[0] # Gradio Interface with Examples gr_interface = gr.Interface( fn=generate_image, inputs=[ gr.Image(type="filepath", label="First Face Image"), gr.Image(type="filepath", label="Second Face Image"), gr.Textbox(label="Prompt"), gr.Dropdown(choices=["Boy", "Girl", "Random"], value="Boy", label="Gender") ], outputs=gr.Image(label="Generated Image"), title="Face Morphing and Image Generation with Stable Diffusion", examples=[ ["yann-lecun.jpg", "isabelle-guyon.jpg", "playing chess", "Boy"] ] ) gr_interface.launch(share=True)