# Prediction interface for Cog ⚙️
# https://github.com/replicate/cog/blob/main/docs/python.md
import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(os.path.join(os.path.dirname(__file__), "../gradio_demo"))
import cv2
import time
import torch
import mimetypes
import subprocess
import numpy as np
from typing import List
from cog import BasePredictor, Input, Path
import PIL
from PIL import Image
import diffusers
from diffusers import LCMScheduler
from diffusers.utils import load_image
from diffusers.models import ControlNetModel
from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
from model_util import get_torch_device
from insightface.app import FaceAnalysis
from transformers import CLIPImageProcessor
from controlnet_util import openpose, get_depth_map, get_canny_image
from diffusers.pipelines.stable_diffusion.safety_checker import (
StableDiffusionSafetyChecker,
)
from pipeline_stable_diffusion_xl_instantid_full import (
StableDiffusionXLInstantIDPipeline,
draw_kps,
)
mimetypes.add_type("image/webp", ".webp")
# GPU global variables
DEVICE = get_torch_device()
DTYPE = torch.float16 if str(DEVICE).__contains__("cuda") else torch.float32
# for `ip-adapter`, `ControlNetModel`, and `stable-diffusion-xl-base-1.0`
CHECKPOINTS_CACHE = "./checkpoints"
CHECKPOINTS_URL = "https://weights.replicate.delivery/default/InstantID/checkpoints.tar"
# for `models/antelopev2`
MODELS_CACHE = "./models"
MODELS_URL = "https://weights.replicate.delivery/default/InstantID/models.tar"
# for the safety checker
SAFETY_CACHE = "./safety-cache"
FEATURE_EXTRACTOR = "./feature-extractor"
SAFETY_URL = "https://weights.replicate.delivery/default/playgroundai/safety-cache.tar"
SDXL_NAME_TO_PATHLIKE = {
# These are all huggingface models that we host via gcp + pget
"stable-diffusion-xl-base-1.0": {
"slug": "stabilityai/stable-diffusion-xl-base-1.0",
"url": "https://weights.replicate.delivery/default/InstantID/models--stabilityai--stable-diffusion-xl-base-1.0.tar",
"path": "checkpoints/models--stabilityai--stable-diffusion-xl-base-1.0",
},
"afrodite-xl-v2": {
"slug": "stablediffusionapi/afrodite-xl-v2",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--afrodite-xl-v2.tar",
"path": "checkpoints/models--stablediffusionapi--afrodite-xl-v2",
},
"albedobase-xl-20": {
"slug": "stablediffusionapi/albedobase-xl-20",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--albedobase-xl-20.tar",
"path": "checkpoints/models--stablediffusionapi--albedobase-xl-20",
},
"albedobase-xl-v13": {
"slug": "stablediffusionapi/albedobase-xl-v13",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--albedobase-xl-v13.tar",
"path": "checkpoints/models--stablediffusionapi--albedobase-xl-v13",
},
"animagine-xl-30": {
"slug": "stablediffusionapi/animagine-xl-30",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--animagine-xl-30.tar",
"path": "checkpoints/models--stablediffusionapi--animagine-xl-30",
},
"anime-art-diffusion-xl": {
"slug": "stablediffusionapi/anime-art-diffusion-xl",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--anime-art-diffusion-xl.tar",
"path": "checkpoints/models--stablediffusionapi--anime-art-diffusion-xl",
},
"anime-illust-diffusion-xl": {
"slug": "stablediffusionapi/anime-illust-diffusion-xl",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--anime-illust-diffusion-xl.tar",
"path": "checkpoints/models--stablediffusionapi--anime-illust-diffusion-xl",
},
"dreamshaper-xl": {
"slug": "stablediffusionapi/dreamshaper-xl",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--dreamshaper-xl.tar",
"path": "checkpoints/models--stablediffusionapi--dreamshaper-xl",
},
"dynavision-xl-v0610": {
"slug": "stablediffusionapi/dynavision-xl-v0610",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--dynavision-xl-v0610.tar",
"path": "checkpoints/models--stablediffusionapi--dynavision-xl-v0610",
},
"guofeng4-xl": {
"slug": "stablediffusionapi/guofeng4-xl",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--guofeng4-xl.tar",
"path": "checkpoints/models--stablediffusionapi--guofeng4-xl",
},
"juggernaut-xl-v8": {
"slug": "stablediffusionapi/juggernaut-xl-v8",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--juggernaut-xl-v8.tar",
"path": "checkpoints/models--stablediffusionapi--juggernaut-xl-v8",
},
"nightvision-xl-0791": {
"slug": "stablediffusionapi/nightvision-xl-0791",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--nightvision-xl-0791.tar",
"path": "checkpoints/models--stablediffusionapi--nightvision-xl-0791",
},
"omnigen-xl": {
"slug": "stablediffusionapi/omnigen-xl",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--omnigen-xl.tar",
"path": "checkpoints/models--stablediffusionapi--omnigen-xl",
},
"pony-diffusion-v6-xl": {
"slug": "stablediffusionapi/pony-diffusion-v6-xl",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--pony-diffusion-v6-xl.tar",
"path": "checkpoints/models--stablediffusionapi--pony-diffusion-v6-xl",
},
"protovision-xl-high-fidel": {
"slug": "stablediffusionapi/protovision-xl-high-fidel",
"url": "https://weights.replicate.delivery/default/InstantID/models--stablediffusionapi--protovision-xl-high-fidel.tar",
"path": "checkpoints/models--stablediffusionapi--protovision-xl-high-fidel",
},
"RealVisXL_V3.0_Turbo": {
"slug": "SG161222/RealVisXL_V3.0_Turbo",
"url": "https://weights.replicate.delivery/default/InstantID/models--SG161222--RealVisXL_V3.0_Turbo.tar",
"path": "checkpoints/models--SG161222--RealVisXL_V3.0_Turbo",
},
"RealVisXL_V4.0_Lightning": {
"slug": "SG161222/RealVisXL_V4.0_Lightning",
"url": "https://weights.replicate.delivery/default/InstantID/models--SG161222--RealVisXL_V4.0_Lightning.tar",
"path": "checkpoints/models--SG161222--RealVisXL_V4.0_Lightning",
},
}
def convert_from_cv2_to_image(img: np.ndarray) -> Image:
return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
def convert_from_image_to_cv2(img: Image) -> np.ndarray:
return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
def resize_img(
input_image,
max_side=1280,
min_side=1024,
size=None,
pad_to_max_side=False,
mode=PIL.Image.BILINEAR,
base_pixel_number=64,
):
w, h = input_image.size
if size is not None:
w_resize_new, h_resize_new = size
else:
ratio = min_side / min(h, w)
w, h = round(ratio * w), round(ratio * h)
ratio = max_side / max(h, w)
input_image = input_image.resize([round(ratio * w), round(ratio * h)], mode)
w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
input_image = input_image.resize([w_resize_new, h_resize_new], mode)
if pad_to_max_side:
res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
offset_x = (max_side - w_resize_new) // 2
offset_y = (max_side - h_resize_new) // 2
res[offset_y : offset_y + h_resize_new, offset_x : offset_x + w_resize_new] = (
np.array(input_image)
)
input_image = Image.fromarray(res)
return input_image
def download_weights(url, dest):
start = time.time()
print("[!] Initiating download from URL: ", url)
print("[~] Destination path: ", dest)
command = ["pget", "-vf", url, dest]
if ".tar" in url:
command.append("-x")
try:
subprocess.check_call(command, close_fds=False)
except subprocess.CalledProcessError as e:
print(
f"[ERROR] Failed to download weights. Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}."
)
raise
print("[+] Download completed in: ", time.time() - start, "seconds")
class Predictor(BasePredictor):
def setup(self) -> None:
"""Load the model into memory to make running multiple predictions efficient"""
if not os.path.exists(CHECKPOINTS_CACHE):
download_weights(CHECKPOINTS_URL, CHECKPOINTS_CACHE)
if not os.path.exists(MODELS_CACHE):
download_weights(MODELS_URL, MODELS_CACHE)
self.face_detection_input_width, self.face_detection_input_height = 640, 640
self.app = FaceAnalysis(
name="antelopev2",
root="./",
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
self.app.prepare(ctx_id=0, det_size=(self.face_detection_input_width, self.face_detection_input_height))
# Path to InstantID models
self.face_adapter = f"./checkpoints/ip-adapter.bin"
controlnet_path = f"./checkpoints/ControlNetModel"
# Load pipeline face ControlNetModel
self.controlnet_identitynet = ControlNetModel.from_pretrained(
controlnet_path,
torch_dtype=DTYPE,
cache_dir=CHECKPOINTS_CACHE,
local_files_only=True,
)
self.setup_extra_controlnets()
self.load_weights("stable-diffusion-xl-base-1.0")
self.setup_safety_checker()
def setup_safety_checker(self):
print(f"[~] Seting up safety checker")
if not os.path.exists(SAFETY_CACHE):
download_weights(SAFETY_URL, SAFETY_CACHE)
self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(
SAFETY_CACHE,
torch_dtype=DTYPE,
local_files_only=True,
)
self.safety_checker.to(DEVICE)
self.feature_extractor = CLIPImageProcessor.from_pretrained(FEATURE_EXTRACTOR)
def run_safety_checker(self, image):
safety_checker_input = self.feature_extractor(image, return_tensors="pt").to(
DEVICE
)
np_image = np.array(image)
image, has_nsfw_concept = self.safety_checker(
images=[np_image],
clip_input=safety_checker_input.pixel_values.to(DTYPE),
)
return image, has_nsfw_concept
def load_weights(self, sdxl_weights):
self.base_weights = sdxl_weights
weights_info = SDXL_NAME_TO_PATHLIKE[self.base_weights]
download_url = weights_info["url"]
path_to_weights_dir = weights_info["path"]
if not os.path.exists(path_to_weights_dir):
download_weights(download_url, path_to_weights_dir)
is_hugging_face_model = "slug" in weights_info.keys()
path_to_weights_file = os.path.join(
path_to_weights_dir,
weights_info.get("file", ""),
)
print(f"[~] Loading new SDXL weights: {path_to_weights_file}")
if is_hugging_face_model:
self.pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
weights_info["slug"],
controlnet=[self.controlnet_identitynet],
torch_dtype=DTYPE,
cache_dir=CHECKPOINTS_CACHE,
local_files_only=True,
safety_checker=None,
feature_extractor=None,
)
self.pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(
self.pipe.scheduler.config
)
else: # e.g. .safetensors, NOTE: This functionality is not being used right now
self.pipe.from_single_file(
path_to_weights_file,
controlnet=self.controlnet_identitynet,
torch_dtype=DTYPE,
cache_dir=CHECKPOINTS_CACHE,
)
self.pipe.load_ip_adapter_instantid(self.face_adapter)
self.setup_lcm_lora()
self.pipe.cuda()
def setup_lcm_lora(self):
print(f"[~] Seting up LCM (just in case)")
lcm_lora_key = "models--latent-consistency--lcm-lora-sdxl"
lcm_lora_path = f"checkpoints/{lcm_lora_key}"
if not os.path.exists(lcm_lora_path):
download_weights(
f"https://weights.replicate.delivery/default/InstantID/{lcm_lora_key}.tar",
lcm_lora_path,
)
self.pipe.load_lora_weights(
"latent-consistency/lcm-lora-sdxl",
cache_dir=CHECKPOINTS_CACHE,
local_files_only=True,
weight_name="pytorch_lora_weights.safetensors",
)
self.pipe.disable_lora()
def setup_extra_controlnets(self):
print(f"[~] Seting up pose, canny, depth ControlNets")
controlnet_pose_model = "thibaud/controlnet-openpose-sdxl-1.0"
controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
controlnet_depth_model = "diffusers/controlnet-depth-sdxl-1.0-small"
for controlnet_key in [
"models--diffusers--controlnet-canny-sdxl-1.0",
"models--diffusers--controlnet-depth-sdxl-1.0-small",
"models--thibaud--controlnet-openpose-sdxl-1.0",
]:
controlnet_path = f"checkpoints/{controlnet_key}"
if not os.path.exists(controlnet_path):
download_weights(
f"https://weights.replicate.delivery/default/InstantID/{controlnet_key}.tar",
controlnet_path,
)
controlnet_pose = ControlNetModel.from_pretrained(
controlnet_pose_model,
torch_dtype=DTYPE,
cache_dir=CHECKPOINTS_CACHE,
local_files_only=True,
).to(DEVICE)
controlnet_canny = ControlNetModel.from_pretrained(
controlnet_canny_model,
torch_dtype=DTYPE,
cache_dir=CHECKPOINTS_CACHE,
local_files_only=True,
).to(DEVICE)
controlnet_depth = ControlNetModel.from_pretrained(
controlnet_depth_model,
torch_dtype=DTYPE,
cache_dir=CHECKPOINTS_CACHE,
local_files_only=True,
).to(DEVICE)
self.controlnet_map = {
"pose": controlnet_pose,
"canny": controlnet_canny,
"depth": controlnet_depth,
}
self.controlnet_map_fn = {
"pose": openpose,
"canny": get_canny_image,
"depth": get_depth_map,
}
def generate_image(
self,
face_image_path,
pose_image_path,
prompt,
negative_prompt,
num_steps,
identitynet_strength_ratio,
adapter_strength_ratio,
pose_strength,
canny_strength,
depth_strength,
controlnet_selection,
guidance_scale,
seed,
scheduler,
enable_LCM,
enhance_face_region,
num_images_per_prompt,
):
if enable_LCM:
self.pipe.enable_lora()
self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
else:
self.pipe.disable_lora()
scheduler_class_name = scheduler.split("-")[0]
add_kwargs = {}
if len(scheduler.split("-")) > 1:
add_kwargs["use_karras_sigmas"] = True
if len(scheduler.split("-")) > 2:
add_kwargs["algorithm_type"] = "sde-dpmsolver++"
scheduler = getattr(diffusers, scheduler_class_name)
self.pipe.scheduler = scheduler.from_config(
self.pipe.scheduler.config,
**add_kwargs,
)
if face_image_path is None:
raise Exception(
f"Cannot find any input face `image`! Please upload the face `image`"
)
face_image = load_image(face_image_path)
face_image = resize_img(face_image)
face_image_cv2 = convert_from_image_to_cv2(face_image)
height, width, _ = face_image_cv2.shape
# Extract face features
face_info = self.app.get(face_image_cv2)
if len(face_info) == 0:
raise Exception(
"Face detector could not find a face in the `image`. Please use a different `image` as input."
)
face_info = sorted(
face_info,
key=lambda x: (x["bbox"][2] - x["bbox"][0]) * x["bbox"][3] - x["bbox"][1],
)[
-1
] # only use the maximum face
face_emb = face_info["embedding"]
face_kps = draw_kps(convert_from_cv2_to_image(face_image_cv2), face_info["kps"])
img_controlnet = face_image
if pose_image_path is not None:
pose_image = load_image(pose_image_path)
pose_image = resize_img(pose_image, max_side=1024)
img_controlnet = pose_image
pose_image_cv2 = convert_from_image_to_cv2(pose_image)
face_info = self.app.get(pose_image_cv2)
if len(face_info) == 0:
raise Exception(
"Face detector could not find a face in the `pose_image`. Please use a different `pose_image` as input."
)
face_info = face_info[-1]
face_kps = draw_kps(pose_image, face_info["kps"])
width, height = face_kps.size
if enhance_face_region:
control_mask = np.zeros([height, width, 3])
x1, y1, x2, y2 = face_info["bbox"]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
control_mask[y1:y2, x1:x2] = 255
control_mask = Image.fromarray(control_mask.astype(np.uint8))
else:
control_mask = None
if len(controlnet_selection) > 0:
controlnet_scales = {
"pose": pose_strength,
"canny": canny_strength,
"depth": depth_strength,
}
self.pipe.controlnet = MultiControlNetModel(
[self.controlnet_identitynet]
+ [self.controlnet_map[s] for s in controlnet_selection]
)
control_scales = [float(identitynet_strength_ratio)] + [
controlnet_scales[s] for s in controlnet_selection
]
control_images = [face_kps] + [
self.controlnet_map_fn[s](img_controlnet).resize((width, height))
for s in controlnet_selection
]
else:
self.pipe.controlnet = self.controlnet_identitynet
control_scales = float(identitynet_strength_ratio)
control_images = face_kps
generator = torch.Generator(device=DEVICE).manual_seed(seed)
print("Start inference...")
print(f"[Debug] Prompt: {prompt}, \n[Debug] Neg Prompt: {negative_prompt}")
self.pipe.set_ip_adapter_scale(adapter_strength_ratio)
images = self.pipe(
prompt=prompt,
negative_prompt=negative_prompt,
image_embeds=face_emb,
image=control_images,
control_mask=control_mask,
controlnet_conditioning_scale=control_scales,
num_inference_steps=num_steps,
guidance_scale=guidance_scale,
height=height,
width=width,
generator=generator,
num_images_per_prompt=num_images_per_prompt,
).images
return images
def predict(
self,
image: Path = Input(
description="Input face image",
),
pose_image: Path = Input(
description="(Optional) reference pose image",
default=None,
),
prompt: str = Input(
description="Input prompt",
default="a person",
),
negative_prompt: str = Input(
description="Input Negative Prompt",
default="",
),
sdxl_weights: str = Input(
description="Pick which base weights you want to use",
default="stable-diffusion-xl-base-1.0",
choices=[
"stable-diffusion-xl-base-1.0",
"juggernaut-xl-v8",
"afrodite-xl-v2",
"albedobase-xl-20",
"albedobase-xl-v13",
"animagine-xl-30",
"anime-art-diffusion-xl",
"anime-illust-diffusion-xl",
"dreamshaper-xl",
"dynavision-xl-v0610",
"guofeng4-xl",
"nightvision-xl-0791",
"omnigen-xl",
"pony-diffusion-v6-xl",
"protovision-xl-high-fidel",
"RealVisXL_V3.0_Turbo",
"RealVisXL_V4.0_Lightning",
],
),
face_detection_input_width: int = Input(
description="Width of the input image for face detection",
default=640,
ge=640,
le=4096,
),
face_detection_input_height: int = Input(
description="Height of the input image for face detection",
default=640,
ge=640,
le=4096,
),
scheduler: str = Input(
description="Scheduler",
choices=[
"DEISMultistepScheduler",
"HeunDiscreteScheduler",
"EulerDiscreteScheduler",
"DPMSolverMultistepScheduler",
"DPMSolverMultistepScheduler-Karras",
"DPMSolverMultistepScheduler-Karras-SDE",
],
default="EulerDiscreteScheduler",
),
num_inference_steps: int = Input(
description="Number of denoising steps",
default=30,
ge=1,
le=500,
),
guidance_scale: float = Input(
description="Scale for classifier-free guidance",
default=7.5,
ge=1,
le=50,
),
ip_adapter_scale: float = Input(
description="Scale for image adapter strength (for detail)", # adapter_strength_ratio
default=0.8,
ge=0,
le=1.5,
),
controlnet_conditioning_scale: float = Input(
description="Scale for IdentityNet strength (for fidelity)", # identitynet_strength_ratio
default=0.8,
ge=0,
le=1.5,
),
enable_pose_controlnet: bool = Input(
description="Enable Openpose ControlNet, overrides strength if set to false",
default=True,
),
pose_strength: float = Input(
description="Openpose ControlNet strength, effective only if `enable_pose_controlnet` is true",
default=0.4,
ge=0,
le=1,
),
enable_canny_controlnet: bool = Input(
description="Enable Canny ControlNet, overrides strength if set to false",
default=False,
),
canny_strength: float = Input(
description="Canny ControlNet strength, effective only if `enable_canny_controlnet` is true",
default=0.3,
ge=0,
le=1,
),
enable_depth_controlnet: bool = Input(
description="Enable Depth ControlNet, overrides strength if set to false",
default=False,
),
depth_strength: float = Input(
description="Depth ControlNet strength, effective only if `enable_depth_controlnet` is true",
default=0.5,
ge=0,
le=1,
),
enable_lcm: bool = Input(
description="Enable Fast Inference with LCM (Latent Consistency Models) - speeds up inference steps, trade-off is the quality of the generated image. Performs better with close-up portrait face images",
default=False,
),
lcm_num_inference_steps: int = Input(
description="Only used when `enable_lcm` is set to True, Number of denoising steps when using LCM",
default=5,
ge=1,
le=10,
),
lcm_guidance_scale: float = Input(
description="Only used when `enable_lcm` is set to True, Scale for classifier-free guidance when using LCM",
default=1.5,
ge=1,
le=20,
),
enhance_nonface_region: bool = Input(
description="Enhance non-face region", default=True
),
output_format: str = Input(
description="Format of the output images",
choices=["webp", "jpg", "png"],
default="webp",
),
output_quality: int = Input(
description="Quality of the output images, from 0 to 100. 100 is best quality, 0 is lowest quality.",
default=80,
ge=0,
le=100,
),
seed: int = Input(
description="Random seed. Leave blank to randomize the seed",
default=None,
),
num_outputs: int = Input(
description="Number of images to output",
default=1,
ge=1,
le=8,
),
disable_safety_checker: bool = Input(
description="Disable safety checker for generated images",
default=False,
),
) -> List[Path]:
"""Run a single prediction on the model"""
# If no seed is provided, generate a random seed
if seed is None:
seed = int.from_bytes(os.urandom(2), "big")
print(f"Using seed: {seed}")
# Load the weights if they are different from the base weights
if sdxl_weights != self.base_weights:
self.load_weights(sdxl_weights)
# Resize the output if the provided dimensions are different from the current ones
if self.face_detection_input_width != face_detection_input_width or self.face_detection_input_height != face_detection_input_height:
print(f"[!] Resizing output to {face_detection_input_width}x{face_detection_input_height}")
self.face_detection_input_width = face_detection_input_width
self.face_detection_input_height = face_detection_input_height
self.app.prepare(ctx_id=0, det_size=(self.face_detection_input_width, self.face_detection_input_height))
# Set up ControlNet selection and their respective strength values (if any)
controlnet_selection = []
if pose_strength > 0 and enable_pose_controlnet:
controlnet_selection.append("pose")
if canny_strength > 0 and enable_canny_controlnet:
controlnet_selection.append("canny")
if depth_strength > 0 and enable_depth_controlnet:
controlnet_selection.append("depth")
# Switch to LCM inference steps and guidance scale if LCM is enabled
if enable_lcm:
num_inference_steps = lcm_num_inference_steps
guidance_scale = lcm_guidance_scale
# Generate
images = self.generate_image(
face_image_path=str(image),
pose_image_path=str(pose_image) if pose_image else None,
prompt=prompt,
negative_prompt=negative_prompt,
num_steps=num_inference_steps,
identitynet_strength_ratio=controlnet_conditioning_scale,
adapter_strength_ratio=ip_adapter_scale,
pose_strength=pose_strength,
canny_strength=canny_strength,
depth_strength=depth_strength,
controlnet_selection=controlnet_selection,
scheduler=scheduler,
guidance_scale=guidance_scale,
seed=seed,
enable_LCM=enable_lcm,
enhance_face_region=enhance_nonface_region,
num_images_per_prompt=num_outputs,
)
# Save the generated images and check for NSFW content
output_paths = []
for i, output_image in enumerate(images):
if not disable_safety_checker:
_, has_nsfw_content_list = self.run_safety_checker(output_image)
has_nsfw_content = any(has_nsfw_content_list)
print(f"NSFW content detected: {has_nsfw_content}")
if has_nsfw_content:
raise Exception(
"NSFW content detected. Try running it again, or try a different prompt."
)
extension = output_format.lower()
extension = "jpeg" if extension == "jpg" else extension
output_path = f"/tmp/out_{i}.{extension}"
print(f"[~] Saving to {output_path}...")
print(f"[~] Output format: {extension.upper()}")
if output_format != "png":
print(f"[~] Output quality: {output_quality}")
save_params = {"format": extension.upper()}
if output_format != "png":
save_params["quality"] = output_quality
save_params["optimize"] = True
output_image.save(output_path, **save_params)
output_paths.append(Path(output_path))
return output_paths